From b6320023c8e34b9710dbe17b3d547a6af7c558c8 Mon Sep 17 00:00:00 2001
From: Rafael Ravedutti <rafaelravedutti@gmail.com>
Date: Mon, 13 Nov 2023 12:53:01 +0100
Subject: [PATCH] Use apply for reductions
Signed-off-by: Rafael Ravedutti <rafaelravedutti@gmail.com>
---
examples/dem.py | 7 ++--
src/pairs/ir/apply.py | 53 ++++++++++++++++++++++++++++++
src/pairs/ir/mutator.py | 10 ++++++
src/pairs/ir/variables.py | 3 +-
src/pairs/mapping/funcs.py | 5 +++
src/pairs/mapping/keywords.py | 10 ++++++
src/pairs/sim/interaction.py | 21 ++++++++++++
src/pairs/sim/simulation.py | 14 ++++++++
src/pairs/transformations/lower.py | 3 ++
9 files changed, 121 insertions(+), 5 deletions(-)
create mode 100644 src/pairs/ir/apply.py
diff --git a/examples/dem.py b/examples/dem.py
index 0ab64a3..ecd4889 100644
--- a/examples/dem.py
+++ b/examples/dem.py
@@ -73,10 +73,9 @@ def linear_spring_dashpot(i, j):
fTabs = min(fTLS_len, f_friction_abs)
fT = fTabs * t
partial_force = fN + fT
- partial_torque = cross(contact_point(i, j) - position[i], partial_force)
- force[i] += partial_force
- torque[i] += partial_torque
+ apply(force, partial_force)
+ apply(torque, cross(contact_point(i, j) - position[i], partial_force))
def euler(i):
@@ -124,7 +123,7 @@ minDiameter_SI = diameter_SI * 0.9
maxDiameter_SI = diameter_SI * 1.1
linkedCellWidth = 1.01 * maxDiameter_SI
-skin = 0.003
+skin = 0.0
ntypes = 1
lnDryResCoeff = math.log(restitutionCoefficient);
diff --git a/src/pairs/ir/apply.py b/src/pairs/ir/apply.py
new file mode 100644
index 0000000..e2658e7
--- /dev/null
+++ b/src/pairs/ir/apply.py
@@ -0,0 +1,53 @@
+from pairs.ir.assign import Assign
+from pairs.ir.ast_node import ASTNode
+from pairs.ir.block import pairs_inline
+from pairs.ir.branches import Filter
+from pairs.ir.lit import Lit
+from pairs.ir.properties import Property
+from pairs.ir.scalars import ScalarOp
+from pairs.ir.vectors import Vector
+from pairs.ir.types import Types
+from pairs.sim.lowerable import FinalLowerable, Lowerable
+
+
+class Apply(Lowerable):
+ def __init__(self, sim, prop, expr, j):
+ assert isinstance(prop, Property), "Apply(): Destination must of Property type."
+ assert prop.type() == expr.type(), "Apply(): Property and expression must be of same type."
+ assert sim.current_apply_list() is not None, "Apply(): Not used within particle interaction."
+ super().__init__(sim)
+ self._prop = prop
+ self._expr = Lit.cvt(sim, expr)
+ self._j = j
+ self._reduction_variable = None
+ sim.current_apply_list().add(self)
+ sim.add_statement(self)
+
+ def __str__(self):
+ return f"Apply<{self._prop, self._expr}>"
+
+ def prop(self):
+ return self._prop
+
+ def expression(self):
+ return self._expr
+
+ def add_reduction_variable(self):
+ self._reduction_variable = self.sim.add_temp_var([0.0, 0.0, 0.0])
+
+ def reduction_variable(self):
+ return self._reduction_variable
+
+ def children(self):
+ return [self._prop, self._expr, self._reduction_variable, self._j]
+
+ @pairs_inline
+ def lower(self):
+ if self.sim._compute_half:
+ for _ in Filter(self.sim,
+ ScalarOp.and_op(j < self.sim.nlocal,
+ ScalarOp.cmp(self.sim.particle_flags[j] & Flags.Fixed, 0))):
+
+ Assign(self.sim, self._prop[self._j], self._prop[self._j] - self._expr)
+
+ Assign(self.sim, self._reduction_variable, self._reduction_variable + self._expr)
diff --git a/src/pairs/ir/mutator.py b/src/pairs/ir/mutator.py
index 8c7cfd6..bd63fbf 100644
--- a/src/pairs/ir/mutator.py
+++ b/src/pairs/ir/mutator.py
@@ -39,6 +39,16 @@ class Mutator:
return ast_node
+ def mutate_Apply(self, ast_node):
+ ast_node._prop = self.mutate(ast_node._prop)
+ ast_node._expr = self.mutate(ast_node._expr)
+ ast_node._j = self.mutate(ast_node._j)
+
+ if ast_node._reduction_variable is not None:
+ ast_node._reduction_variable = self.mutate(ast_node._reduction_variable)
+
+ return ast_node
+
def mutate_ArrayAccess(self, ast_node):
ast_node.array = self.mutate(ast_node.array)
ast_node.partial_indexes = [self.mutate(i) for i in ast_node.partial_indexes]
diff --git a/src/pairs/ir/variables.py b/src/pairs/ir/variables.py
index d57e081..4c1d279 100644
--- a/src/pairs/ir/variables.py
+++ b/src/pairs/ir/variables.py
@@ -3,6 +3,7 @@ from pairs.ir.ast_term import ASTTerm
from pairs.ir.assign import Assign
from pairs.ir.scalars import ScalarOp
from pairs.ir.lit import Lit
+from pairs.ir.operator_class import OperatorClass
class Variables:
@@ -39,7 +40,7 @@ class Variables:
class Var(ASTTerm):
def __init__(self, sim, var_name, var_type, init_value=0, temp=False):
- super().__init__(sim, ScalarOp)
+ super().__init__(sim, OperatorClass.from_type(var_type))
self.var_name = var_name
self.var_type = var_type
self.var_init_value = Lit.cvt(sim, init_value)
diff --git a/src/pairs/mapping/funcs.py b/src/pairs/mapping/funcs.py
index b1689c3..573ee6a 100644
--- a/src/pairs/mapping/funcs.py
+++ b/src/pairs/mapping/funcs.py
@@ -150,6 +150,9 @@ class BuildParticleIR(ast.NodeVisitor):
args = [self.visit(a) for a in node.args]
if self.keywords.exists(func):
+ if func == 'apply':
+ args += [self.ctx_symbols['__j__']]
+
return self.keywords(func, args)
if func in ['delta', 'squared_distance', 'penetration_depth', 'contact_point', 'contact_normal']:
@@ -274,6 +277,8 @@ def compute(sim, func, cutoff_radius=None, symbols={}):
ir.add_symbols({
params[0]: interaction_data.i(),
params[1]: interaction_data.j(),
+ '__i__': interaction_data.i(),
+ '__j__': interaction_data.j(),
'__delta__': interaction_data.delta(),
'__squared_distance__': interaction_data.squared_distance(),
'__penetration_depth__': interaction_data.penetration_depth(),
diff --git a/src/pairs/mapping/keywords.py b/src/pairs/mapping/keywords.py
index 2d7917f..9ae93a6 100644
--- a/src/pairs/mapping/keywords.py
+++ b/src/pairs/mapping/keywords.py
@@ -1,9 +1,11 @@
+from pairs.ir.apply import Apply
from pairs.ir.block import Block
from pairs.ir.branches import Filter
from pairs.ir.lit import Lit
from pairs.ir.loops import Continue
from pairs.ir.math import Abs, Cos, Sin, Sqrt
from pairs.ir.matrices import Matrix
+from pairs.ir.properties import Property
from pairs.ir.quaternions import Quaternion
from pairs.ir.scalars import ScalarOp
from pairs.ir.select import Select
@@ -247,3 +249,11 @@ class Keywords:
2.0 * (quat[1] * quat[3] - quat[0] * quat[2]),
2.0 * (quat[2] * quat[3] + quat[0] * quat[1]),
1.0 - 2.0 * quat[1] * quat[1] - 2.0 * quat[2] * quat[2] ])
+
+ def keyword_apply(self, args):
+ assert len(args) == 3, "apply() keyword requires two parameters."
+ prop = args[0]
+ expr = args[1]
+ assert isinstance(prop, Property), "apply(): First argument must be a property."
+ assert prop.type() == expr.type(), "apply(): Property and expression must be of same type."
+ Apply(self.sim, prop, expr, args[2])
diff --git a/src/pairs/sim/interaction.py b/src/pairs/sim/interaction.py
index 02bf67f..3fcf52f 100644
--- a/src/pairs/sim/interaction.py
+++ b/src/pairs/sim/interaction.py
@@ -1,3 +1,4 @@
+from pairs.ir.assign import Assign
from pairs.ir.ast_term import ASTTerm
from pairs.ir.scalars import ScalarOp
from pairs.ir.block import Block, pairs_inline
@@ -136,6 +137,7 @@ class ParticleInteraction(Lowerable):
self.interactions_data = []
self.blocks = [Block(sim, []) for _ in range(sim.max_shapes() * self.ncases)]
self.active_block = None
+ self.apply_list = set()
def add_statement(self, stmt):
self.active_block.add_statement(stmt)
@@ -143,6 +145,7 @@ class ParticleInteraction(Lowerable):
def __iter__(self):
self.sim.add_statement(self)
self.sim.enter(self)
+ self.sim.use_apply_list(self.apply_list)
# Neighbors vary across iterations
for shape in range(self.sim.max_shapes()):
@@ -151,6 +154,7 @@ class ParticleInteraction(Lowerable):
self.interactions_data.append(InteractionData(self.sim, shape))
yield self.interactions_data[-1]
+ self.sim.release_apply_list()
self.sim.leave()
self.active_block = None
@@ -163,6 +167,9 @@ class ParticleInteraction(Lowerable):
for i in ParticleFor(self.sim):
for _ in Filter(self.sim, ScalarOp.cmp(self.sim.particle_flags[i] & Flags.Fixed, 0)):
+ for app in self.apply_list:
+ app.add_reduction_variable()
+
interaction = 0
for neigh in NeighborFor(self.sim, i, cell_lists, neighbor_lists):
interaction_data = self.interactions_data[interaction]
@@ -227,5 +234,19 @@ class ParticleInteraction(Lowerable):
self.sim.add_statement(Filter(self.sim, cutoff_condition, self.blocks[interaction]))
interaction += 1
+ prop_reductions = {}
+ for app in self.apply_list:
+ prop = app.prop()
+ reduction = app.reduction_variable()
+
+ if prop not in prop_reductions:
+ prop_reductions[prop] = reduction
+
+ else:
+ prop_reductions[prop] = prop_reductions[prop] + reduction
+
+ for prop, reduction in prop_reductions.items():
+ Assign(self.sim, prop[i], reduction)
+
else:
raise Exception("Interactions among more than two particles are currently not supported.")
diff --git a/src/pairs/sim/simulation.py b/src/pairs/sim/simulation.py
index 0f13067..cff4475 100644
--- a/src/pairs/sim/simulation.py
+++ b/src/pairs/sim/simulation.py
@@ -80,6 +80,11 @@ class Simulation:
self._use_contact_history = use_contact_history
self._contact_history = ContactHistory(self) if use_contact_history else None
self._shapes = shapes
+ self._compute_half = False
+ self._apply_list = None
+
+ def compute_half(self):
+ self._compute_half = True
def use_double_precision(self):
return self._double_prec
@@ -281,6 +286,15 @@ class Simulation:
else:
self.nested_count += 1
+ def use_apply_list(self, apply_list):
+ self._apply_list = apply_list
+
+ def release_apply_list(self):
+ self._apply_list = None
+
+ def current_apply_list(self):
+ return self._apply_list
+
def vtk_output(self, filename, frequency=0):
self.vtk_file = filename
self.vtk_frequency = frequency
diff --git a/src/pairs/transformations/lower.py b/src/pairs/transformations/lower.py
index 487a78a..9033499 100644
--- a/src/pairs/transformations/lower.py
+++ b/src/pairs/transformations/lower.py
@@ -11,6 +11,9 @@ class Lower(Mutator):
def set_data(self, data):
self.lower_finals = data[0]
+ def mutate_Apply(self, ast_node):
+ return self.mutate_Unknown(ast_node)
+
def mutate_Unknown(self, ast_node):
if isinstance(ast_node, Lowerable) or (self.lower_finals and isinstance(ast_node, FinalLowerable)):
self.lowered_nodes += 1
--
GitLab