Skip to content

Commit a3c2426

Browse files
mlouboutmloubout
committed
compiler: add rudimentary support for multi-cond buffering
1 parent 1c65e48 commit a3c2426

5 files changed

Lines changed: 153 additions & 88 deletions

File tree

‎devito/ir/clusters/algorithms.py‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -254,6 +254,7 @@ def guard(clusters):
254254

255255
# Chain together all `cds` conditions from all expressions in `c`
256256
guards = {}
257+
mode = sympy.Or
257258
for cd in cds:
258259
# `BOTTOM` parent implies a guard that lives outside of
259260
# any iteration space, which corresponds to the placeholder None
@@ -270,6 +271,7 @@ def guard(clusters):
270271

271272
# Pull `cd` from any expr
272273
condition = guards.setdefault(k, [])
274+
mode = mode and cd.relation
273275
for e in exprs:
274276
try:
275277
condition.append(e.conditionals[cd])
@@ -284,7 +286,8 @@ def guard(clusters):
284286
conditionals.pop(cd, None)
285287
exprs[i] = e.func(*e.args, conditionals=conditionals)
286288

287-
guards = {d: sympy.And(*v, evaluate=False) for d, v in guards.items()}
289+
# Combination mode is And by default and Or if all conditions are
290+
guards = {d: mode(*v, evaluate=False) for d, v in guards.items()}
288291

289292
# Construct a guarded Cluster
290293
processed.append(c.rebuild(exprs=exprs, guards=guards))

‎devito/ir/equations/equation.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -223,7 +223,7 @@ def __new__(cls, *args, **kwargs):
223223
else:
224224
cond = diff2sympy(lower_exprs(d.condition))
225225
if d._factor is not None:
226-
cond = sympy.And(cond, GuardFactor(d))
226+
cond = d.relation(cond, GuardFactor(d))
227227
conditionals[d] = cond
228228
# Replace dimension with index
229229
index = d.index

‎devito/types/dimension.py‎

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -860,6 +860,8 @@ class ConditionalDimension(DerivedDimension):
860860
If True, use `self`, rather than the parent Dimension, to
861861
index into arrays. A typical use case is when arrays are accessed
862862
indirectly via the ``condition`` expression.
863+
relation: Or/And, default=And
864+
How this ConditionalDimension will be combined with other ones.
863865
864866
Examples
865867
--------
@@ -913,10 +915,10 @@ class ConditionalDimension(DerivedDimension):
913915
is_Conditional = True
914916

915917
__rkwargs__ = DerivedDimension.__rkwargs__ + \
916-
('factor', 'condition', 'indirect')
918+
('factor', 'condition', 'indirect', 'relation')
917919

918920
def __init_finalize__(self, name, parent=None, factor=None, condition=None,
919-
indirect=False, **kwargs):
921+
indirect=False, relation=sympy.And, **kwargs):
920922
# `parent=None` degenerates to a ConditionalDimension outside of
921923
# any iteration space
922924
if parent is None:
@@ -937,6 +939,7 @@ def __init_finalize__(self, name, parent=None, factor=None, condition=None,
937939

938940
self._condition = condition
939941
self._indirect = indirect
942+
self._relation = relation
940943

941944
@property
942945
def uses_symbolic_factor(self):
@@ -978,6 +981,10 @@ def condition(self):
978981
def indirect(self):
979982
return self._indirect
980983

984+
@property
985+
def relation(self):
986+
return self._relation
987+
981988
@cached_property
982989
def free_symbols(self):
983990
retval = set(super().free_symbols)

‎examples/userapi/05_conditional_dimension.ipynb‎

Lines changed: 104 additions & 83 deletions
Original file line numberDiff line numberDiff line change
@@ -71,6 +71,7 @@
7171
"name": "stderr",
7272
"output_type": "stream",
7373
"text": [
74+
"NUMA domain count autodetection failed, assuming 1\n",
7475
"Operator `Kernel` ran in 0.01 s\n"
7576
]
7677
},
@@ -158,75 +159,77 @@
158159
"output_type": "stream",
159160
"text": [
160161
"\n",
161-
" Symbol defining a non-convex iteration sub-space derived from a ``parent``\n",
162-
" Dimension, implemented by the compiler generating conditional \"if-then\" code\n",
163-
" within the parent Dimension's iteration space.\n",
162+
"Symbol defining a non-convex iteration sub-space derived from a ``parent``\n",
163+
"Dimension, implemented by the compiler generating conditional \"if-then\" code\n",
164+
"within the parent Dimension's iteration space.\n",
164165
"\n",
165-
" Parameters\n",
166-
" ----------\n",
167-
" name : str\n",
168-
" Name of the dimension.\n",
169-
" parent : Dimension\n",
170-
" The parent Dimension.\n",
171-
" factor : int, optional, default=None\n",
172-
" The number of iterations between two executions of the if-branch. If None\n",
173-
" (default), ``condition`` must be provided.\n",
174-
" condition : expr-like, optional, default=None\n",
175-
" An arbitrary SymPy expression, typically involving the ``parent``\n",
176-
" Dimension. When it evaluates to True, the if-branch is executed. If None\n",
177-
" (default), ``factor`` must be provided.\n",
178-
" indirect : bool, optional, default=False\n",
179-
" If True, use `self`, rather than the parent Dimension, to\n",
180-
" index into arrays. A typical use case is when arrays are accessed\n",
181-
" indirectly via the ``condition`` expression.\n",
166+
"Parameters\n",
167+
"----------\n",
168+
"name : str\n",
169+
" Name of the dimension.\n",
170+
"parent : Dimension\n",
171+
" The parent Dimension.\n",
172+
"factor : int, optional, default=None\n",
173+
" The number of iterations between two executions of the if-branch. If None\n",
174+
" (default), ``condition`` must be provided.\n",
175+
"condition : expr-like, optional, default=None\n",
176+
" An arbitrary SymPy expression, typically involving the ``parent``\n",
177+
" Dimension. When it evaluates to True, the if-branch is executed. If None\n",
178+
" (default), ``factor`` must be provided.\n",
179+
"indirect : bool, optional, default=False\n",
180+
" If True, use `self`, rather than the parent Dimension, to\n",
181+
" index into arrays. A typical use case is when arrays are accessed\n",
182+
" indirectly via the ``condition`` expression.\n",
183+
"relation: Or/And, default=And\n",
184+
" How this ConditionalDimension will be combined with other ones.\n",
182185
"\n",
183-
" Examples\n",
184-
" --------\n",
185-
" Among the other things, ConditionalDimensions are indicated to implement\n",
186-
" Function subsampling. In the following example, an Operator evaluates the\n",
187-
" Function ``g`` and saves its content into ``f`` every ``factor=4`` iterations.\n",
186+
"Examples\n",
187+
"--------\n",
188+
"Among the other things, ConditionalDimensions are indicated to implement\n",
189+
"Function subsampling. In the following example, an Operator evaluates the\n",
190+
"Function ``g`` and saves its content into ``f`` every ``factor=4`` iterations.\n",
188191
"\n",
189-
" >>> from devito import Dimension, ConditionalDimension, Function, Eq, Operator\n",
190-
" >>> size, factor = 16, 4\n",
191-
" >>> i = Dimension(name='i')\n",
192-
" >>> ci = ConditionalDimension(name='ci', parent=i, factor=factor)\n",
193-
" >>> g = Function(name='g', shape=(size,), dimensions=(i,))\n",
194-
" >>> f = Function(name='f', shape=(int(size/factor),), dimensions=(ci,))\n",
195-
" >>> op = Operator([Eq(g, 1), Eq(f, g)])\n",
192+
">>> from devito import Dimension, ConditionalDimension, Function, Eq, Operator\n",
193+
">>> size, factor = 16, 4\n",
194+
">>> i = Dimension(name='i')\n",
195+
">>> ci = ConditionalDimension(name='ci', parent=i, factor=factor)\n",
196+
">>> g = Function(name='g', shape=(size,), dimensions=(i,))\n",
197+
">>> f = Function(name='f', shape=(int(size/factor),), dimensions=(ci,))\n",
198+
">>> op = Operator([Eq(g, 1), Eq(f, g)])\n",
196199
"\n",
197-
" The Operator generates the following for-loop (pseudocode)\n",
200+
"The Operator generates the following for-loop (pseudocode)\n",
198201
"\n",
199-
" .. code-block:: C\n",
202+
".. code-block:: C\n",
200203
"\n",
201-
" for (int i = i_m; i <= i_M; i += 1) {\n",
202-
" g[i] = 1;\n",
203-
" if (i%4 == 0) {\n",
204-
" f[i / 4] = g[i];\n",
205-
" }\n",
206-
" }\n",
204+
" for (int i = i_m; i <= i_M; i += 1) {\n",
205+
" g[i] = 1;\n",
206+
" if (i%4 == 0) {\n",
207+
" f[i / 4] = g[i];\n",
208+
" }\n",
209+
" }\n",
207210
"\n",
208-
" Another typical use case is when one needs to constrain the execution of\n",
209-
" loop iterations so that certain conditions are honoured. The following\n",
210-
" artificial example uses ConditionalDimension to guard against out-of-bounds\n",
211-
" accesses in indirectly accessed arrays.\n",
211+
"Another typical use case is when one needs to constrain the execution of\n",
212+
"loop iterations so that certain conditions are honoured. The following\n",
213+
"artificial example uses ConditionalDimension to guard against out-of-bounds\n",
214+
"accesses in indirectly accessed arrays.\n",
212215
"\n",
213-
" >>> from sympy import And\n",
214-
" >>> ci = ConditionalDimension(name='ci', parent=i,\n",
215-
" ... condition=And(g[i] > 0, g[i] < 4, evaluate=False))\n",
216-
" >>> f = Function(name='f', shape=(int(size/factor),), dimensions=(ci,))\n",
217-
" >>> op = Operator(Eq(f[g[i]], f[g[i]] + 1))\n",
216+
">>> from sympy import And\n",
217+
">>> ci = ConditionalDimension(name='ci', parent=i,\n",
218+
"... condition=And(g[i] > 0, g[i] < 4, evaluate=False))\n",
219+
">>> f = Function(name='f', shape=(int(size/factor),), dimensions=(ci,))\n",
220+
">>> op = Operator(Eq(f[g[i]], f[g[i]] + 1))\n",
218221
"\n",
219-
" The Operator generates the following for-loop (pseudocode)\n",
222+
"The Operator generates the following for-loop (pseudocode)\n",
220223
"\n",
221-
" .. code-block:: C\n",
224+
".. code-block:: C\n",
222225
"\n",
223-
" for (int i = i_m; i <= i_M; i += 1) {\n",
224-
" if (g[i] > 0 && g[i] < 4) {\n",
225-
" f[g[i]] = f[g[i]] + 1;\n",
226-
" }\n",
227-
" }\n",
226+
" for (int i = i_m; i <= i_M; i += 1) {\n",
227+
" if (g[i] > 0 && g[i] < 4) {\n",
228+
" f[g[i]] = f[g[i]] + 1;\n",
229+
" }\n",
230+
" }\n",
228231
"\n",
229-
" \n"
232+
"\n"
230233
]
231234
}
232235
],
@@ -248,9 +251,7 @@
248251
{
249252
"cell_type": "code",
250253
"execution_count": 5,
251-
"metadata": {
252-
"scrolled": true
253-
},
254+
"metadata": {},
254255
"outputs": [
255256
{
256257
"name": "stderr",
@@ -321,14 +322,18 @@
321322
"output_type": "stream",
322323
"text": [
323324
"START(section0)\n",
324-
"for (int x = x_m; x <= x_M; x += 1)\n",
325+
"#pragma omp parallel num_threads(nthreads)\n",
325326
"{\n",
326-
" #pragma omp simd aligned(f:32)\n",
327-
" for (int y = y_m; y <= y_M; y += 1)\n",
327+
" #pragma omp for schedule(static,1)\n",
328+
" for (int x = x_m; x <= x_M; x += 1)\n",
328329
" {\n",
329-
" if (f[x + 1][y + 1] > 0)\n",
330+
" #pragma omp simd aligned(f:16)\n",
331+
" for (int y = y_m; y <= y_M; y += 1)\n",
330332
" {\n",
331-
" f[x + 1][y + 1] = f[x + 1][y + 1] + 1;\n",
333+
" if (f[x + 1][y + 1] > 0)\n",
334+
" {\n",
335+
" f[x + 1][y + 1] = f[x + 1][y + 1] + 1;\n",
336+
" }\n",
332337
" }\n",
333338
" }\n",
334339
"}\n",
@@ -397,14 +402,18 @@
397402
"output_type": "stream",
398403
"text": [
399404
"START(section0)\n",
400-
"for (int x = x_m; x <= x_M; x += 1)\n",
405+
"#pragma omp parallel num_threads(nthreads)\n",
401406
"{\n",
402-
" #pragma omp simd aligned(f,g:32)\n",
403-
" for (int y = y_m; y <= y_M; y += 1)\n",
407+
" #pragma omp for schedule(static,1)\n",
408+
" for (int x = x_m; x <= x_M; x += 1)\n",
404409
" {\n",
405-
" if (y < 5 && g[x + 1][y + 1] != 0)\n",
410+
" #pragma omp simd aligned(f,g:16)\n",
411+
" for (int y = y_m; y <= y_M; y += 1)\n",
406412
" {\n",
407-
" f[x + 1][y + 1] = f[x + 1][y + 1] + g[x + 1][y + 1];\n",
413+
" if (y < 5 && g[x + 1][y + 1] != 0)\n",
414+
" {\n",
415+
" f[x + 1][y + 1] = f[x + 1][y + 1] + g[x + 1][y + 1];\n",
416+
" }\n",
408417
" }\n",
409418
" }\n",
410419
"}\n",
@@ -489,13 +498,17 @@
489498
"output_type": "stream",
490499
"text": [
491500
"START(section0)\n",
492-
"for (int x = x_m; x <= x_M; x += 1)\n",
501+
"#pragma omp parallel num_threads(nthreads)\n",
493502
"{\n",
494-
" for (int y = y_m; y <= y_M; y += 1)\n",
503+
" #pragma omp for schedule(static,1)\n",
504+
" for (int x = x_m; x <= x_M; x += 1)\n",
495505
" {\n",
496-
" if (y < 5 && g[x + 1][y + 1] != 0)\n",
506+
" for (int y = y_m; y <= y_M; y += 1)\n",
497507
" {\n",
498-
" h[x + 1][y + 1] = g[x + 1][y + 1] + h[x + 1][y + 1];\n",
508+
" if (y < 5 && g[x + 1][y + 1] != 0)\n",
509+
" {\n",
510+
" h[x + 1][y + 1] = g[x + 1][y + 1] + h[x + 1][y + 1];\n",
511+
" }\n",
499512
" }\n",
500513
" }\n",
501514
"}\n",
@@ -563,11 +576,15 @@
563576
"output_type": "stream",
564577
"text": [
565578
"START(section0)\n",
566-
"for (int i = i_m; i <= i_M; i += 1)\n",
579+
"#pragma omp parallel num_threads(nthreads)\n",
567580
"{\n",
568-
" if ((i)%(cif) == 0)\n",
581+
" #pragma omp for schedule(static,1)\n",
582+
" for (int i = i_m; i <= i_M; i += 1)\n",
569583
" {\n",
570-
" f[i / cif] = g[i];\n",
584+
" if ((i)%(cif) == 0)\n",
585+
" {\n",
586+
" f[i / cif] = g[i];\n",
587+
" }\n",
571588
" }\n",
572589
"}\n",
573590
"STOP(section0,timers)\n",
@@ -678,13 +695,17 @@
678695
"output_type": "stream",
679696
"text": [
680697
"START(section0)\n",
681-
"for (int x = x_m; x <= x_M; x += 1)\n",
698+
"#pragma omp parallel num_threads(nthreads)\n",
682699
"{\n",
683-
" for (int y = y_m; y <= y_M; y += 1)\n",
700+
" #pragma omp for collapse(2) schedule(static,1) reduction(+:g[1])\n",
701+
" for (int x = x_m; x <= x_M; x += 1)\n",
684702
" {\n",
685-
" if (f[x][y] != 0)\n",
703+
" for (int y = y_m; y <= y_M; y += 1)\n",
686704
" {\n",
687-
" g[1] += 1;\n",
705+
" if (f[x][y] != 0)\n",
706+
" {\n",
707+
" g[1] += 1;\n",
708+
" }\n",
688709
" }\n",
689710
" }\n",
690711
"}\n",
@@ -694,7 +715,7 @@
694715
{
695716
"data": {
696717
"text/plain": [
697-
"10"
718+
"np.int32(10)"
698719
]
699720
},
700721
"execution_count": 11,
@@ -831,7 +852,7 @@
831852
"name": "python",
832853
"nbconvert_exporter": "python",
833854
"pygments_lexer": "ipython3",
834-
"version": "3.11.2"
855+
"version": "3.13.11"
835856
}
836857
},
837858
"nbformat": 4,

0 commit comments

Comments
 (0)