def add(self, alias, intervals, aliaseds, distances):
assert len(aliaseds) == len(distances)
self[alias] = (intervals, aliaseds, distances)
# Update the index_mapper
for i in intervals:
d = i.dim
if d in self.index_mapper:
continue
elif d.is_Shifted:
self.index_mapper[d.parent] = d
elif d.is_Incr:
# IncrDimensions must be substituted with ShiftedDimensions
# to access the temporaries, otherwise we would go OOB
# E.g., r[xs][ys][z] => `xs/ys` must start at 0, not at `x0_blk0`
# as in the case of blocking
self.index_mapper[d] = ShiftedDimension(d, "%ss" % d.name)
def make_schedule(cluster, aliases, in_writeto, options):
"""
Create a Schedule from an AliasMapper.
"""
max_par = options['cire-maxpar']
dmapper = {}
processed = []
for alias, v in aliases.items():
imapper = {**{i.dim: i for i in v.intervals},
**{i.dim.parent: i for i in v.intervals if i.dim.is_NonlinearDerived}}
intervals = []
writeto = []
sub_iterators = {}
indicess = [[] for _ in v.distances]
for i in cluster.ispace.intervals:
try:
interval = imapper[i.dim]
except KeyError:
# E.g., `x0_blk0` or (`a[y_m+1]` => `y not in imapper`)
intervals.append(i)
continue
assert i.stamp >= interval.stamp
if not (writeto or interval != interval.zero() or in_writeto(i.dim, cluster)):
# The alias doesn't require a temporary Dimension along i.dim
intervals.append(i)
continue
assert not i.dim.is_NonlinearDerived
# `i.dim` is necessarily part of the write-to region, so
# we have to adjust the Interval's stamp. For example, consider
# `i=x[0,0]<1>` and `interval=x[-4,4]<0>`; here we need to
# use `<1>` as stamp, which is what appears in `cluster`
interval = interval.lift(i.stamp)
# We further bump the interval stamp if we were requested to trade
# fusion for more collapse-parallelism
interval = interval.lift(interval.stamp + int(max_par))
writeto.append(interval)
intervals.append(interval)
if i.dim.is_Incr:
# Suitable ShiftedDimensions must be used to avoid OOB accesses.
# E.g., r[xs][ys][z] => both `xs` and `ys` must start at 0,
# not at `x0_blk0`
try:
d = dmapper[i.dim]
except KeyError:
d = dmapper[i.dim] = ShiftedDimension(i.dim, name="%ss" % i.dim.name)
sub_iterators[i.dim] = d
else:
d = i.dim
# Given the iteration `interval`, lower distances to indices
for distance, indices in zip(v.distances, indicess):
indices.append(d - interval.lower + distance[interval.dim])
# The alias write-to space
writeto = IterationSpace(IntervalGroup(writeto), sub_iterators)
# The alias iteration space
intervals = IntervalGroup(intervals, cluster.ispace.relations)
ispace = IterationSpace(intervals, cluster.sub_iterators, cluster.directions)
ispace = ispace.augment(sub_iterators)
processed.append(ScheduledAlias(alias, writeto, ispace, v.aliaseds, indicess))
# Sort by write-to region for deterministic code generation
processed = sorted(processed, key=lambda i: i.writeto)
return Schedule(*processed, dmapper=dmapper)