api: fix handling of multiple conditions for buffering - #2850
Conversation
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #2850 +/- ##
==========================================
+ Coverage 83.64% 83.66% +0.01%
==========================================
Files 257 257
Lines 54235 54524 +289
Branches 4629 4680 +51
==========================================
+ Hits 45367 45618 +251
- Misses 8072 8099 +27
- Partials 796 807 +11
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
| return CondNe(*self.args, evaluate=False) | ||
|
|
||
| @property | ||
| def _as_min(self): |
There was a problem hiding this comment.
I would drop this and rather have a singledispatch handler for CondEq where necessary
| expr = uxreplace(expr, {d: IntDiv(index, d.symbolic_factor)}) | ||
|
|
||
| # Merge conditionals when possible. E.g if we have an implicit_dim | ||
| # and there is a dimension with the same parent, we ca merged |
There was a problem hiding this comment.
Dimension
"ca merged"
"their conditions"
you could also make the example a bit more practical
| for d in input_expr.implicit_dims: | ||
| if d not in conditionals: | ||
| continue | ||
| for cd in dict(conditionals): |
| # Replace the ConditionalDimensions in `expr` | ||
| for d, cond in conditionals.items(): | ||
| # Replace dimension with index | ||
| index = d.index |
There was a problem hiding this comment.
you can spare this line
| ispace = IterationSpace(intervals, iterators) | ||
|
|
||
| # Construct the conditionals and replace the ConditionalDimensions in `expr` | ||
| # Construct the conditionals |
There was a problem hiding this comment.
I think we should place this whole block of code, which constructs/lowers the conditionals, into its own separate functions, and a docstring with some examples
ef708e5 to
b997156
Compare
7a1a6aa to
c7786ea
Compare
|
Check out this pull request on See visual diffs & provide feedback on Jupyter Notebooks. Powered by ReviewNB |
f904760 to
0500469
Compare
| shift = relational_shift(cond, d.parent) | ||
| expr = uxreplace(expr, {d: IntDiv(index, d.symbolic_factor) + shift}) | ||
|
|
||
| # Merge conditionals when possible. E.g if we have an implicit_dim |
There was a problem hiding this comment.
btw this block imho deserves its own function
| if d is not dim: | ||
| continue | ||
|
|
||
| if d in c0.guards and not c0.guards[d].has(Mod): |
There was a problem hiding this comment.
searching for Mod is a bit meh, I'd rather add a special guard to ir/support/guards.py and look for that instead (there's quite a few already in there!)
| _actions_from_update_memcpy(c, d, clusters, actions, sregistry) | ||
| elif d.is_Custom and is_integer(c.ispace[d].size): | ||
| _actions_from_init(c, d, actions) | ||
| _actions_from_init(c, d, clusters, actions) |
|
|
||
|
|
||
| def _actions_from_init(c, d, actions): | ||
| def _actions_from_init(c, d, clusters, actions): |
0500469 to
30790f0
Compare
89aea51 to
99f0a9e
Compare
| if d not in conditionals: | ||
| continue | ||
| for cd in list(conditionals): | ||
| if cd.parent == d.parent and cd is not d: |
There was a problem hiding this comment.
nitpicking: cd.parent is d.parent ...
| continue | ||
| for cd in list(conditionals): | ||
| if cd.parent == d.parent and cd is not d: | ||
| cond = conditionals.pop(d) |
There was a problem hiding this comment.
nitpicking: you don't need this line here, you may just .pop in the else branch, and avoid the = cond in the if branch
|
|
||
| # Replace the ConditionalDimensions in `expr` | ||
| for d, cond in conditionals.items(): | ||
| # Replace dimension with index |
| g = c0.guards.get(d) | ||
| # Explicit compute guards need no pipeline; memcpy clusters | ||
| # still need WithLock for the copy-back sync | ||
| if g is not None and not wraps_memcpy(c0): |
There was a problem hiding this comment.
I think a lot of the complexity inherent in this guard and the nested one could be mitigated by simply introducing somehting along the lines of c0.guards.has(d, GuardFactor)
|
|
||
| # Alias deferred entries to an existing f-buffer; create one if none | ||
| for f, deferred in extras.items(): | ||
| reusable = [mapper[(ff, kk)] for (ff, kk) in mapper if ff == f] |
| # Add all missing buffering dimensions and reorder to | ||
| # avoid duplicates with different ordering | ||
| ispaces.add(c.ispace.insert(self.dim, edims).reorder()) | ||
| elif ((self.f in c.scope.reads or self.f in c.scope.writes) and |
There was a problem hiding this comment.
self.f in c.scope.functions or in fact self.f in c.functions should suffice
|
|
||
| lhs = b.indexify()._subs(v.xd, v.first_idx.b) | ||
| rhs = f.indexify()._subs(v.dim, v.first_idx.f) | ||
| for b, vb in descriptors.items(): |
There was a problem hiding this comment.
another potential occurrence of descriptors.flatten
| ) | ||
| else: | ||
| size = async_degree | ||
| for k, ck in groupby(clusters, key=lambda c: c.guards): |
There was a problem hiding this comment.
ultra-uber-nitpicking, typically we use letter g for groups
| return p, offset | ||
|
|
||
|
|
||
| def _buffer_dim(f, key, ck): |
There was a problem hiding this comment.
this is potentially a legacy over-complication
is it basically what in PRO refer to as .hdim ?
because imho, if so, we could/should somehow abstract the concept here, and just use f.hdim (or whatever name you wanna give it)
99f0a9e to
bebbf4f
Compare
ab8cd52 to
b1147aa
Compare
267a4ed to
cfa4221
Compare
| # the purpose of protecting from OOB accesses | ||
| cds = [d for d in cds if not d.indirect] | ||
| modes = [cd.relation for cd in cds] | ||
| strict = ConditionalDimension._STRICT |
There was a problem hiding this comment.
STRICT instead of _STRICT since it's not a private class attribute
| # still need WithLock for the copy-back sync. An "explicit" | ||
| # guard is a plain relation on `d` (e.g., `d == K`); subsampling | ||
| # guards (GuardFactor) still require the standard async pipeline | ||
| if d in c0.guards and not wraps_memcpy(c0) \ |
There was a problem hiding this comment.
ultra nitpick, for homogeneity we typically do:
if d in c0.guards and \
not wraps_memcpy(c0) and \
not c0.guards.has(d, GuardFactor):
There was a problem hiding this comment.
blank lines surrounding this new block would help also
| descriptors = {b: BufferDescriptor(f, b, bfmap[f]) | ||
| for f, b in self.mapper.items() | ||
| if f in bfmap} | ||
| descriptors = {} |
| properties = c.properties.sequentialize(d) | ||
| if not isinstance(d, BufferDimension): | ||
| if not isinstance(d, BufferDimension) and \ | ||
| not _explicit_guard(c.guards.get(d, None), d): |
There was a problem hiding this comment.
since _explicit_guard is called only here, you can probably move the not isinstance(d, BufferDimension) part of the condition inside of it, thus simplifying this caller site
| index_mapper = {i: mds[(v.xd, i)] for i in v.indices} | ||
| for a in accesses: | ||
| subs[a.access] = b.indexed[[index_mapper.get(i, i) for i in a]] | ||
| subs = defaultdict(dict) |
There was a problem hiding this comment.
I'd do this, it'd shorten this loop body significantly
|
|
||
| cls = callback or Array | ||
| name = sregistry.make_name(prefix=f'{f.name}b') | ||
| return cls(name=name, dimensions=dimensions, dtype=f.dtype, |
There was a problem hiding this comment.
u now might be able to fit this in two lines
|
|
||
| # Indirect accessors define the ispace only for a read-only streamed | ||
| # buffer, where nothing iterates the buffer's own Dimensions directly | ||
| if not ispaces: |
There was a problem hiding this comment.
this could be glued with the if len(ispaces) > 1:, u could have:
if len(ispaces) == 0:
# Indirect accessors define the ispace only for a read-only streamed
# buffer, where nothing iterates the buffer's own Dimensions directly
ispace = indirect.pop()
elif len(ispaces) == 1:
# <comment>
ispace = ispaces.pop()
else:
# Best effort to make buffering work in the presence of multiple
# IterationSpaces
<code>
ispace = ispaces.pop()
return ispace
There was a problem hiding this comment.
in retrospect, do you really need indirect ? are u sure u cannot just reuse the same ispaces ? I guess u know better than me...
There was a problem hiding this comment.
actually can probably just use ispaces yeah, droping indirect
| buf = next((b for b in reusable | ||
| if set(k) & set(b.dimensions)), None) | ||
| if buf is None and reusable: | ||
| buf = reusable[0] |
There was a problem hiding this comment.
are u expecting len(reusable) == 1 ? if yes, is it worth asserting it? if not, why do we pick [0]
There was a problem hiding this comment.
No, that was lazy, tweked for better selection
|
|
||
|
|
||
| def _make_buffer(f, dim, k, ck, xds, async_degree, sregistry, callback): | ||
| """Build (or retrieve) the buffer Array for `f` along `dim` under guards `k`.""" |
There was a problem hiding this comment.
then instead of _make_buffer I'd call it differently (pick_buffer ?) and remove the parentheses surrounding "or retrieve" since it's very important , not a minor detail
| return self.getters[d] | ||
| raise KeyError | ||
|
|
||
| def dindex(self, dim): |
There was a problem hiding this comment.
why not just a def index override? if the arg is an integer you call into super(), otherwise do the logic below
There was a problem hiding this comment.
It's only used in one place (for transforms) and it makes the standard index work better as the dict's default one
No description provided.