|
15 | 15 | from devito.passes.clusters.cse import _cse |
16 | 16 | from devito.passes.clusters.utils import expose_tuning_knobs |
17 | 17 | from devito.symbolics import ( |
18 | | - Uxmapper, estimate_cost, retrieve_functions, reuse_if_untouched, search, sympy_dtype, |
19 | | - uxreplace |
| 18 | + INT, Uxmapper, estimate_cost, retrieve_functions, reuse_if_untouched, search, |
| 19 | + sympy_dtype, uxreplace |
20 | 20 | ) |
21 | 21 | from devito.tools import ( |
22 | 22 | Reconstructable, Stamp, as_mapper, as_tuple, flatten, frozendict, generator, |
@@ -295,6 +295,10 @@ def _do_generate(self, exprs, exclude, cbk_search, cbk_compose=None): |
295 | 295 |
|
296 | 296 | class CireInvariants(CireTransformerLegacy, Queue): |
297 | 297 |
|
| 298 | + # Predicate on Cluster used to pick which ones this pass fires on. |
| 299 | + # Subclasses override to target a different kind of cluster. |
| 300 | + _cluster_filter = staticmethod(lambda c: c.is_dense) |
| 301 | + |
298 | 302 | def __init__(self, sregistry, options, platform): |
299 | 303 | super().__init__(sregistry, options, platform) |
300 | 304 |
|
@@ -324,7 +328,8 @@ def callback(self, clusters, prefix, xtracted=None): |
324 | 328 | key = lambda c: self._lookup_key(c, d) |
325 | 329 | processed = list(clusters) |
326 | 330 | for ak, group in as_mapper(clusters, key=key).items(): |
327 | | - g = [c for c in group if c.is_dense and c not in xtracted] |
| 331 | + g = [c for c in group |
| 332 | + if self._cluster_filter(c) and c not in xtracted] |
328 | 333 | if not g: |
329 | 334 | continue |
330 | 335 |
|
@@ -387,6 +392,52 @@ def _generate(self, cgroup, exclude): |
387 | 392 | yield self._do_generate(exprs, exclude, cbk_search) |
388 | 393 |
|
389 | 394 |
|
| 395 | +def _is_floor(e): |
| 396 | + return getattr(e, 'is_Function', False) and e.func.__name__ == 'floor' |
| 397 | + |
| 398 | + |
| 399 | +class CireInvariantsFloor(CireInvariants): |
| 400 | + |
| 401 | + """ |
| 402 | + Hoist `INT(floor(x))` expressions -- typically the integer cell of a |
| 403 | + sparse position `INT(floor((c - o)/h))` -- into an int32 preamble Array. |
| 404 | + `CireInvariantsElementary`'s commutativity gate rejects both `floor(...)` |
| 405 | + (Cast in the argument makes `is_commutative == False`) and |
| 406 | + `INT(floor(...))` (a `Cast`, `is_commutative is None`), so we handle |
| 407 | + them here. Fires on every cluster: hoisting an invariant integer floor |
| 408 | + out of an inner loop is a strict win regardless of the cluster kind. |
| 409 | + """ |
| 410 | + |
| 411 | + _cluster_filter = staticmethod(lambda c: True) |
| 412 | + |
| 413 | + def _generate(self, cgroup, exclude): |
| 414 | + counter = generator() |
| 415 | + make_i32 = lambda: Symbol(name=f'dummy{counter()}', dtype=np.int32) |
| 416 | + |
| 417 | + mapper = Uxmapper() |
| 418 | + for e in cgroup.exprs: |
| 419 | + for f in search(e, _is_floor, 'all', 'bfs'): |
| 420 | + cand = INT(f) |
| 421 | + if not {a.function for a in cand.free_symbols} & exclude: |
| 422 | + mapper.add(cand, make_i32, None) |
| 423 | + |
| 424 | + yield mapper |
| 425 | + |
| 426 | + def _choose(self, aliases, cgroup, mapper): |
| 427 | + # Skip score-based filtering: hoisting a floor out of an inner |
| 428 | + # loop is a strict flops win regardless of working-set size. |
| 429 | + exprs = cgroup.exprs |
| 430 | + |
| 431 | + aliases = AliasList(aliases) |
| 432 | + if not aliases: |
| 433 | + return exprs, aliases |
| 434 | + |
| 435 | + subs = {k: v for k, v in mapper.items() |
| 436 | + if v.free_symbols & set(aliases.aliaseds)} |
| 437 | + exprs = [uxreplace(e, subs) for e in exprs] |
| 438 | + return exprs, aliases |
| 439 | + |
| 440 | + |
390 | 441 | class CireDerivatives(CireTransformerLegacy): |
391 | 442 |
|
392 | 443 | def __init__(self, sregistry, options, platform): |
@@ -519,7 +570,8 @@ def _cbk_search2(self, expr, rank): |
519 | 570 | # Subpass mapper |
520 | 571 | modes = { |
521 | 572 | 'invariants': [CireInvariantsElementary, |
522 | | - CireInvariantsDivs], |
| 573 | + CireInvariantsDivs, |
| 574 | + CireInvariantsFloor], |
523 | 575 | 'eval-derivs': [CireEvalDerivatives], # NOTE: legacy pass |
524 | 576 | 'index-derivs': [CireIndexDerivatives], |
525 | 577 | } |
@@ -707,8 +759,14 @@ def lower_aliases(aliases, meta, opt_maxpar, opt_block_temps): |
707 | 759 | # not as an Indexed index. Then, it needs to be added to |
708 | 760 | # the `writeto` region too |
709 | 761 | interval = i |
| 762 | + elif writeto: |
| 763 | + # `d` is inner to the writeto region and unused by the |
| 764 | + # alias -- skip it, otherwise the preamble would iterate |
| 765 | + # over it and store the same value repeatedly. |
| 766 | + continue |
710 | 767 | else: |
711 | | - # E.g., `x0_blk0` or (`a[y_m+1]` => `y not in imapper`) |
| 768 | + # E.g., `x0_blk0` (outer to writeto, must stay in ispace |
| 769 | + # so later passes -- `Lift`, `Fuse` -- see it) |
712 | 770 | intervals[d] = i |
713 | 771 | continue |
714 | 772 |
|
@@ -1565,10 +1623,10 @@ def make_rotations_table(d, v): |
1565 | 1623 |
|
1566 | 1624 | def cit(ispace0, ispace1): |
1567 | 1625 | """ |
1568 | | - The Common IterationIntervals of two IterationSpaces. |
| 1626 | + The Common IterationIntervals of two IterationSpaces (the shared prefix). |
1569 | 1627 | """ |
1570 | 1628 | found = [] |
1571 | | - for it0, it1 in zip(ispace0.itintervals, ispace1.itintervals, strict=True): |
| 1629 | + for it0, it1 in zip(ispace0.itintervals, ispace1.itintervals, strict=False): |
1572 | 1630 | if it0 == it1: |
1573 | 1631 | found.append(it0) |
1574 | 1632 | else: |
|
0 commit comments