def __init__(self, function, contracted_dims, accessv, n, async_degree):
self.function = function
self.accessv = accessv
contraction_mapper = {}
index_mapper = {}
dims = list(function.dimensions)
for d in contracted_dims:
assert d in function.dimensions
# Determine the buffer size along `d`
indices = filter_ordered(i.indices[d] for i in accessv.accesses)
slots = [i.xreplace({d: 0, d.spacing: 1}) for i in indices]
size = max(slots) - min(slots) + 1
if async_degree is not None:
if async_degree < size:
warning("Ignoring provided asynchronous degree as it'd be "
"too small for the required buffer (provided %d, "
"but need at least %d for `%s`)"
% (async_degree, size, function.name))
else:
size = async_degree
# Replace `d` with a suitable CustomDimension
bd = CustomDimension('db%d' % n, 0, size-1, size, d)
contraction_mapper[d] = dims[dims.index(d)] = bd
if size > 1:
# Create the necessary SteppingDimensions for indexing
sd = SteppingDimension(name='sb%d' % n, parent=bd)
index_mapper.update({i: i.xreplace({d: sd}) for i in indices})
else:
# Special case, no need to keep a SteppingDimension around
index_mapper.update({i: 0 for i in indices})
self.contraction_mapper = contraction_mapper
self.index_mapper = index_mapper
# Track the SubDimensions used to index into `function`
subdims_mapper = DefaultOrderedDict(set)
for e in accessv.mapper:
try:
# Case 1: implicitly via SubDomains
m = {d.root: v for d, v in e.subdomain.dimension_map.items()}
except AttributeError:
# Case 2: explicitly via the lower-level SubDimension API
m = {i.root: i for i in e.free_symbols
if isinstance(i, Dimension) and (i.is_Sub or not i.is_Derived)}
for d, v in m.items():
subdims_mapper[d].add(v)
if any(len(v) > 1 for v in subdims_mapper.values()):
# Non-uniform SubDimensions. At this point we're going to raise
# an exception. It's either illegal or still unsupported
for v in subdims_mapper.values():
for d0, d1 in combinations(v, 2):
if d0.overlap(d1):
raise InvalidOperator("Cannot apply `buffering` to `%s` as it "
"is accessed over the overlapping "
" SubDimensions `<%s, %s>`" %
(function, d0, d1))
self.subdims_mapper = None
raise NotImplementedError("`buffering` does not support multiple "
"non-overlapping SubDimensions yet.")
else:
self.subdims_mapper = {d: v.pop() for d, v in subdims_mapper.items()}
self.buffer = Array(name='%sb' % function.name,
dimensions=dims,
dtype=function.dtype,
halo=function.halo,
space='mapped')
class Buffer(object):
"""
A buffer with metadata attached.
Parameters
----------
function : DiscreteFunction
The object for which the buffer is created.
contracted_dims : list of Dimension
The Dimensions in `function` to be contracted, that is to be replaced
by ModuloDimensions.
accessv : AccessValue
All accesses involving `function`.
options : dict, optional
The compilation options. See `buffering.__doc__`.
sregistry : SymbolRegistry
The symbol registry, to create unique names for buffers and Dimensions.
bds : dict, optional
All CustomDimensions created to define buffer dimensions, potentially
reusable in the creation of this buffer. The object gets updated if new
CustomDimensions are created.
mds : dict, optional
All ModuloDimensions created to index into other buffers, potentially reusable
for indexing into this buffer. The object gets updated if new ModuloDimensions
are created.
"""
def __init__(self, function, contracted_dims, accessv, options, sregistry,
bds=None, mds=None):
# Parse compilation options
async_degree = options['buf-async-degree']
space = options['buf-mem-space']
dtype = options['buf-dtype'](function)
self.function = function
self.accessv = accessv
self.contraction_mapper = {}
self.index_mapper = defaultdict(dict)
self.sub_iterators = defaultdict(list)
self.subdims_mapper = DefaultOrderedDict(set)
# Create the necessary ModuloDimensions for indexing into the buffer
# E.g., `u[time,x] + u[time+1,x] -> `ub[sb0,x] + ub[sb1,x]`, where `sb0`
# and `sb1` are ModuloDimensions starting at `time` and `time+1` respectively
dims = list(function.dimensions)
for d in contracted_dims:
assert d in function.dimensions
# Determine the buffer size, and therefore the span of the ModuloDimension,
# along the contracting Dimension `d`
indices = filter_ordered(i.indices[d] for i in accessv.accesses)
slots = [i.subs({d: 0, d.spacing: 1}) for i in indices]
try:
size = max(slots) - min(slots) + 1
except TypeError:
# E.g., special case `slots=[-1 + time/factor, 2 + time/factor]`
# Resort to the fast vector-based comparison machinery (rather than
# the slower sympy.simplify)
slots = [Vector(i) for i in slots]
size = int((vmax(*slots) - vmin(*slots) + 1)[0])
if async_degree is not None:
if async_degree < size:
warning("Ignoring provided asynchronous degree as it'd be "
"too small for the required buffer (provided %d, "
"but need at least %d for `%s`)"
% (async_degree, size, function.name))
else:
size = async_degree
# Replace `d` with a suitable CustomDimension `bd`
name = sregistry.make_name(prefix='db')
bd = bds.setdefault((d, size), CustomDimension(name, 0, size-1, size, d))
self.contraction_mapper[d] = dims[dims.index(d)] = bd
# Finally create the ModuloDimensions as children of `bd`
if size > 1:
# Note: indices are sorted so that the semantic order (sb0, sb1, sb2)
# follows SymPy's index ordering (time, time-1, time+1) after modulo
# replacement, so that associativity errors are consistent. This very
# same strategy is also applied in clusters/algorithms/Stepper
p, _ = offset_from_centre(d, indices)
indices = sorted(indices,
key=lambda i: -np.inf if i - p == 0 else (i - p))
for i in indices:
name = sregistry.make_name(prefix='sb')
md = mds.setdefault((bd, i), ModuloDimension(name, bd, i, size))
self.index_mapper[d][i] = md
self.sub_iterators[d.root].append(md)
else:
assert len(indices) == 1
self.index_mapper[d][indices[0]] = 0
# Track the SubDimensions used to index into `function`
for e in accessv.mapper:
m = {i.root: i for i in e.free_symbols
if isinstance(i, Dimension) and (i.is_Sub or not i.is_Derived)}
for d, v in m.items():
self.subdims_mapper[d].add(v)
if any(len(v) > 1 for v in self.subdims_mapper.values()):
# Non-uniform SubDimensions. At this point we're going to raise
# an exception. It's either illegal or still unsupported
for v in self.subdims_mapper.values():
for d0, d1 in combinations(v, 2):
if d0.overlap(d1):
raise InvalidOperator("Cannot apply `buffering` to `%s` as it "
"is accessed over the overlapping "
" SubDimensions `<%s, %s>`" %
(function, d0, d1))
raise NotImplementedError("`buffering` does not support multiple "
"non-overlapping SubDimensions yet.")
else:
self.subdims_mapper = {d: v.pop() for d, v in self.subdims_mapper.items()}
# Build and sanity-check the buffer IterationIntervals
self.itintervals_mapper = {}
for e in accessv.mapper:
for i in e.ispace.itintervals:
v = self.itintervals_mapper.setdefault(i.dim, i.args)
if v != self.itintervals_mapper[i.dim]:
raise NotImplementedError("Cannot apply `buffering` as the buffered "
"function `%s` is accessed over multiple, "
"non-compatible iteration spaces along the "
"Dimension `%s`" % (function.name, i.dim))
# Also add IterationIntervals for initialization along `x`, should `xi` be
# the only written Dimension in the `x` hierarchy
for d, (interval, _, _) in list(self.itintervals_mapper.items()):
for i in d._defines:
self.itintervals_mapper.setdefault(i, (interval.relaxed, (), Forward))
# Finally create the actual buffer
self.buffer = Array(name=sregistry.make_name(prefix='%sb' % function.name),
dimensions=dims,
dtype=dtype,
halo=function.halo,
space=space)
def __repr__(self):
return "Buffer[%s,<%s>]" % (self.buffer.name,
','.join(str(i) for i in self.contraction_mapper))
@property
def size(self):
return np.prod([v.symbolic_size for v in self.contraction_mapper.values()])
@property
def is_read(self):
return self.accessv.is_read
@property
def is_readonly(self):
return self.is_read and self.lastwrite is None
@property
def firstread(self):
return self.accessv.firstread
@property
def lastwrite(self):
return self.accessv.lastwrite
@property
def has_uniform_subdims(self):
return self.subdims_mapper is not None
@cached_property
def indexed(self):
return self.buffer.indexed
@cached_property
def index_mapper_flat(self):
ret = {}
for mapper in self.index_mapper.values():
ret.update(mapper)
return ret
@cached_property
def writeto(self):
"""
The `writeto` IterationSpace, that is the iteration space that must be
iterated over in order to initialize the buffer.
"""
intervals = []
sub_iterators = {}
directions = {}
for d in self.buffer.dimensions:
try:
interval, si, direction = self.itintervals_mapper[d]
except KeyError:
# E.g., the contraction Dimension `db0`
assert d in self.contraction_mapper.values()
interval, si, direction = Interval(d, 0, 0), (), Forward
intervals.append(interval)
sub_iterators[d] = si
directions[d] = direction
relations = (self.buffer.dimensions,)
intervals = IntervalGroup(intervals, relations=relations)
return IterationSpace(intervals, sub_iterators, directions)
@cached_property
def written(self):
"""
The `written` IterationSpace, that is the iteration space that must be
iterated over in order to read all of the written buffer values.
"""
intervals = []
sub_iterators = {}
directions = {}
for dd in self.function.dimensions:
d = dd.xreplace(self.subdims_mapper)
try:
interval, si, direction = self.itintervals_mapper[d]
except KeyError:
# E.g., d=time_sub
assert d.is_NonlinearDerived
d = d.root
interval, si, direction = self.itintervals_mapper[d]
intervals.append(interval)
sub_iterators[d] = si + as_tuple(self.sub_iterators[d])
directions[d] = direction
relations = (tuple(i.dim for i in intervals),)
intervals = IntervalGroup(intervals, relations=relations)
return IterationSpace(intervals, sub_iterators, directions)
@cached_property
def lastmap(self):
"""
A mapper from contracted Dimensions to a 2-tuple of indices representing,
respectively, the "last" write to the buffer and the "last" read from the
buffered Function. For example, `{time: (sb1, time+1)}`.
"""
mapper = {}
for d, m in self.index_mapper.items():
try:
func = max if self.written.directions[d.root] is Forward else min
v = func(m)
except TypeError:
func = vmax if self.written.directions[d.root] is Forward else vmin
v = func(*[Vector(i) for i in m])[0]
mapper[d] = Map(m[v], v)
return mapper
@cached_property
def initmap(self):
"""
A mapper from contracted Dimensions to indices representing the min points
for buffer initialization. For example, in the case of a forward-propagating
`time` Dimension, we could have `{time: (time_m + db0) % 2, (time_m + db0)}`;
likewise, for backwards, `{time: (time_M - 2 + db0) % 4, time_M - 2 + db0}`.
"""
mapper = {}
for d, bd in self.contraction_mapper.items():
indices = list(self.index_mapper[d])
# The buffer is initialized at `d_m(d_M) - offset`. E.g., a buffer with
# six slots, used to replace a buffered Function accessed at `d-3`, `d`
# and `d + 2`, will have `offset = 3`
p, offset = offset_from_centre(d, indices)
if self.written.directions[d.root] is Forward:
v = p.subs(d.root, d.root.symbolic_min) - offset + bd
else:
v = p.subs(d.root, d.root.symbolic_max) - offset + bd
mapper[d] = Map(v % bd.symbolic_size, v)
return mapper
