|
| 1 | +from devito.ir.clusters import Queue |
| 2 | +from devito.ir.support import GuardFactor |
| 3 | +from devito.symbolics import IntDiv, uxreplace |
| 4 | +from devito.tools import timed_pass |
| 5 | + |
| 6 | +__all__ = ['subsampling'] |
| 7 | + |
| 8 | + |
| 9 | +@timed_pass() |
| 10 | +def subsampling(clusters): |
| 11 | + """ |
| 12 | + Strengthen factor-based ConditionalDimension guards into strided iteration. |
| 13 | +
|
| 14 | + A Function subsampled by `factor` along `d` is written under a |
| 15 | + `d % factor == 0` guard, so the enclosing loop visits every point of `d` |
| 16 | + and discards all but one in `factor` of them. Where legal, drop the guard |
| 17 | + and step the loop `factor` points at a time instead |
| 18 | +
|
| 19 | + for (x = x_m; x <= x_M; x += 1) for (x = ...; x <= x_M; x += xf) |
| 20 | + if (x%xf == 0) -> g[x/xf] = ... |
| 21 | + g[x/xf] = ... |
| 22 | +
|
| 23 | +The loop is taken over the subsampled Dimension `cd` itself, with a *unit* |
| 24 | + stride, and every access to the parent `d` is rescaled to `cd*factor`: |
| 25 | +
|
| 26 | + for (x = x_m; x <= x_M; x += 1) for (xs = xs_m; xs <= xs_M; xs++) |
| 27 | + if (x%xf == 0) -> g[xs] = u[xs*xf] |
| 28 | + g[x/xf] = u[x] |
| 29 | +
|
| 30 | + A non-unitary increment is deliberately avoided: it would escape the |
| 31 | + Lamport-style dependence analysis the compiler relies on. The loop bounds |
| 32 | + are `cd`'s own runtime values, which `ConditionalDimension._arg_values` |
| 33 | + derives as `ceil(x_m/factor) .. floor(x_M/factor)`, so no anchoring of the |
| 34 | + loop minimum is needed either. |
| 35 | +
|
| 36 | + Parameters |
| 37 | + ---------- |
| 38 | + clusters : tuple of Clusters |
| 39 | + Input Clusters, subject of the optimization pass. |
| 40 | + """ |
| 41 | + return Subsampling().process(clusters) |
| 42 | + |
| 43 | + |
| 44 | +class Subsampling(Queue): |
| 45 | + |
| 46 | + """ |
| 47 | + Replace uniform factor-based guards with strided iteration. |
| 48 | +
|
| 49 | + Clusters are grouped by iteration space prefix alone, deliberately *not* |
| 50 | + by guard: a guard may only become a stride if every Cluster sharing the |
| 51 | + prefix carries the very same one. Striding a loop nest that also hosts an |
| 52 | + unguarded expression would run that expression on a subset of its |
| 53 | + iterations -- and dropping the guard is what makes the nest shared in the |
| 54 | + first place, since `stree_build` reuses a subtree across Clusters with |
| 55 | + matching guards (see `reuse_whole_subtree`). |
| 56 | + """ |
| 57 | + |
| 58 | + def callback(self, clusters, prefix): |
| 59 | + if not prefix: |
| 60 | + return clusters |
| 61 | + |
| 62 | + d = prefix[-1].dim |
| 63 | + |
| 64 | + guards = {c.guards.get(d) for c in clusters} |
| 65 | + if len(guards) > 1: |
| 66 | + # Non-uniformly guarded, hence `d`'s loop nest is shared with |
| 67 | + # expressions that must keep running at every point |
| 68 | + return clusters |
| 69 | + |
| 70 | + guard = guards.pop() |
| 71 | + if not all(_is_strideable(c, d, guard) for c in clusters): |
| 72 | + return clusters |
| 73 | + |
| 74 | + cd = guard.d |
| 75 | + factor = cd.symbolic_factor |
| 76 | + |
| 77 | + processed = [] |
| 78 | + for c in clusters: |
| 79 | + # Iterate the subsampled Dimension itself, with unit stride, and |
| 80 | + # rescale every parent access `d` to `cd*factor` |
| 81 | + ispace = c.ispace.switch(d, cd) |
| 82 | + |
| 83 | + # `IntDiv(d, factor)` is how the subsampled access was lowered; |
| 84 | + # with `d -> cd*factor` it collapses to `cd`, so rewrite it |
| 85 | + # directly rather than leaving `(cd*factor)/factor` behind |
| 86 | + subs = {IntDiv(d, factor): cd, d: cd*factor} |
| 87 | + exprs = [uxreplace(e, subs) for e in c.exprs] |
| 88 | + |
| 89 | + # `cd` inherits `d`'s computational properties: it covers the very |
| 90 | + # same points, just addressed in subsampled coordinates. Without |
| 91 | + # this the loop would lose PARALLEL and so all of its OpenMP |
| 92 | + properties = c.properties.drop(d) |
| 93 | + properties = properties.add(cd, c.properties[d]) |
| 94 | + |
| 95 | + processed.append(c.rebuild(exprs=exprs, ispace=ispace, |
| 96 | + guards=c.guards.popany(d), |
| 97 | + properties=properties)) |
| 98 | + |
| 99 | + return processed |
| 100 | + |
| 101 | + |
| 102 | +def _is_strideable(cluster, d, guard): |
| 103 | + """ |
| 104 | + True if iteration over `d` may step `guard`'s subsampling factor points at |
| 105 | + a time in place of evaluating `guard` at every point. |
| 106 | + """ |
| 107 | + # The guard must be exactly a subsampling guard. Anything `And`/`Or`-ed in |
| 108 | + # (e.g. a user-supplied `condition`) must still be evaluated pointwise |
| 109 | + if not isinstance(guard, GuardFactor): |
| 110 | + return False |
| 111 | + |
| 112 | + if d not in cluster.ispace.itdims: |
| 113 | + return False |
| 114 | + |
| 115 | + # Skipping points is only sound if no dependence is carried along `d`; |
| 116 | + # otherwise the discarded iterations may be doing work that `d`'s |
| 117 | + # remaining iterations rely upon |
| 118 | + if not cluster.properties.is_parallel_relaxed(d): |
| 119 | + return False |
| 120 | + |
| 121 | + # Non-affine accesses may read the points that would be skipped |
| 122 | + return cluster.properties.is_affine(d) |
0 commit comments