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