def st_make_halo(stree): """ Add :class:`NodeHalo`s to a :class:`ScheduleTree`. A HaloNode captures the halo exchanges that should take place before executing the sub-tree; these are described by means of a :class:`HaloScheme`. """ # Build a HaloScheme for each expression bundle halo_schemes = {} for n in findall(stree, lambda i: i.is_Exprs): try: halo_schemes[n] = HaloScheme(n.exprs, n.ispace, n.dspace) except HaloSchemeException as e: if configuration['mpi']: raise RuntimeError(str(e)) # Insert the HaloScheme at a suitable level in the ScheduleTree mapper = {} for k, hs in halo_schemes.items(): for f, v in hs.fmapper.items(): spot = k ancestors = [n for n in k.ancestors if n.is_Iteration] for n in ancestors: test0 = any(n.dim is i.dim for i in v.halos) test1 = n.dim not in [i.root for i in v.loc_indices] if test0 or test1: spot = n break mapper.setdefault(spot, []).append((f, v)) for spot, entries in mapper.items(): insert(NodeHalo(HaloScheme(fmapper=dict(entries))), spot.parent, [spot]) return stree
def st_make_halo(stree): """ Add NodeHalos to a ScheduleTree. A NodeHalo captures the halo exchanges that should take place before executing the sub-tree; these are described by means of a HaloScheme. """ # Build a HaloScheme for each expression bundle halo_schemes = {} for n in findall(stree, lambda i: i.is_Exprs): try: halo_schemes[n] = HaloScheme(n.exprs, n.ispace) except HaloSchemeException as e: if configuration['mpi']: raise RuntimeError(str(e)) # Insert the HaloScheme at a suitable level in the ScheduleTree mapper = {} for k, hs in halo_schemes.items(): for f, v in hs.fmapper.items(): spot = k ancestors = [n for n in k.ancestors if n.is_Iteration] for n in ancestors: test0 = any(n.dim is i.dim for i in v.halos) test1 = n.dim not in [i.root for i in v.loc_indices] if test0 or test1: spot = n break mapper.setdefault(spot, []).append((f, v)) for spot, entries in mapper.items(): insert(NodeHalo(HaloScheme(fmapper=dict(entries))), spot.parent, [spot]) return stree
def stree_section(stree): """ Add NodeSections to a ScheduleTree. A NodeSection, or simply "section", defines a sub-tree with the following properties: * The root is a node of type NodeSection; * The immediate children of the root are nodes of type NodeIteration; * The Dimensions of the immediate children are either: * identical, OR * different, but all of type SubDimension; * The Dimension of the immediate children cannot be a TimeDimension. """ class Section(object): def __init__(self, node): self.parent = node.parent try: self.dim = node.dim except AttributeError: self.dim = None self.nodes = [node] def is_compatible(self, node): return self.parent == node.parent and self.dim.root == node.dim.root # Search candidate sections sections = [] for i in range(stree.height): # Find all sections at depth `i` section = None for n in findall(stree, filter_=lambda n: n.depth == i): if any(p in flatten(s.nodes for s in sections) for p in n.ancestors): # Already within a section continue elif n.is_Sync: # SyncNodes are self-contained sections.append(Section(n)) section = None elif n.is_Iteration: if n.dim.is_Time and SEQUENTIAL in n.properties: # If n.dim.is_Time, we end up here in 99.9% of the cases. # Sometimes, however, time is a PARALLEL Dimension (e.g., # think of `norm` Operators) section = None elif section is None or not section.is_compatible(n): section = Section(n) sections.append(section) else: section.nodes.append(n) else: section = None # Transform the schedule tree by adding in sections for i in sections: insert(NodeSection(), i.parent, i.nodes) return stree
def st_section(stree): """ Add :class:`NodeSection` to a :class:`ScheduleTree`. A section defines a sub-tree with the following properties: :: * The root is a node of type :class:`NodeSection`; * The immediate children of the root are nodes of type :class:`NodeIteration` and have same parent. * The :class:`Dimension` of the immediate children are either: :: * identical, OR * different, but all of type :class:`SubDimension`; * The :class:`Dimension` of the immediate children cannot be a :class:`TimeDimension`. """ class Section(object): def __init__(self, node): self.parent = node.parent self.dim = node.dim self.nodes = [node] def is_compatible(self, node): return (self.parent == node.parent and (self.dim == node.dim or node.dim.is_Sub)) # Search candidate sections sections = [] for i in range(stree.height): # Find all sections at depth `i` section = None for n in findall(stree, filter_=lambda n: n.depth == i): if any(p in flatten(s.nodes for s in sections) for p in n.ancestors): # Already within a section continue elif not n.is_Iteration or n.dim.is_Time: section = None elif section is None or not section.is_compatible(n): section = Section(n) sections.append(section) else: section.nodes.append(n) # Transform the schedule tree by adding in sections for i in sections: insert(NodeSection(), i.parent, i.nodes) return stree
def st_section(stree): """ Add NodeSections to a ScheduleTree. A NodeSection, or simply "section", defines a sub-tree with the following properties: * The root is a node of type NodeSection; * The immediate children of the root are nodes of type NodeIteration; * The Dimensions of the immediate children are either: * identical, OR * different, but all of type SubDimension; * The Dimension of the immediate children cannot be a TimeDimension. """ class Section(object): def __init__(self, node): self.parent = node.parent self.dim = node.dim self.nodes = [node] def is_compatible(self, node): return self.parent == node.parent and self.dim.root == node.dim.root # Search candidate sections sections = [] for i in range(stree.height): # Find all sections at depth `i` section = None for n in findall(stree, filter_=lambda n: n.depth == i): if any(p in flatten(s.nodes for s in sections) for p in n.ancestors): # Already within a section continue elif not n.is_Iteration or n.dim.is_Time: section = None elif section is None or not section.is_compatible(n): section = Section(n) sections.append(section) else: section.nodes.append(n) # Transform the schedule tree by adding in sections for i in sections: insert(NodeSection(), i.parent, i.nodes) return stree
def stree_make_halo(stree): """ Add NodeHalos to a ScheduleTree. A NodeHalo captures the halo exchanges that should take place before executing the sub-tree; these are described by means of a HaloScheme. """ # Build a HaloScheme for each expression bundle halo_schemes = {} for n in findall(stree, lambda i: i.is_Exprs): try: halo_schemes[n] = HaloScheme(n.exprs, n.ispace) except HaloSchemeException as e: if configuration['mpi']: raise RuntimeError(str(e)) # Split a HaloScheme based on where it should be inserted # For example, it's possible that, for a given HaloScheme, a Function's # halo needs to be exchanged at a certain `stree` depth, while another # Function's halo needs to be exchanged before some other nodes mapper = {} for k, hs in halo_schemes.items(): for f, v in hs.fmapper.items(): spot = k ancestors = [n for n in k.ancestors if n.is_Iteration] for n in ancestors: # Place the halo exchange right before the first # distributed Dimension which requires it if any(i.dim in n.dim._defines for i in v.halos): spot = n break mapper.setdefault(spot, []).append(hs.project(f)) # Now fuse the HaloSchemes at the same `stree` depth and perform the insertion for spot, halo_schemes in mapper.items(): insert(NodeHalo(HaloScheme.union(halo_schemes)), spot.parent, [spot]) return stree
def st_make_halo(stree): """ Add :class:`NodeHalo` to a :class:`ScheduleTree`. A halo node describes what halo exchanges should take place before executing the sub-tree. """ if not configuration['mpi']: # TODO: This will be dropped as soon as stronger analysis will have # been implemented return stree processed = {} for n in LevelOrderIter(stree, stop=lambda i: i.parent in processed): if not n.is_Iteration: continue exprs = flatten(i.exprs for i in findall(n, lambda i: i.is_Exprs)) try: halo_scheme = HaloScheme(exprs) if n.dim in halo_scheme.dmapper: processed[n] = NodeHalo(halo_scheme) except HaloSchemeException: # We should get here only when trying to compute a halo # scheme for a group of expressions that belong to different # iteration spaces. We expect proper halo schemes to be built # as the `stree` visit proceeds. # TODO: However, at the end, we should check that a halo scheme, # possibly even a "void" one, has been built for *all* of the # expressions, and error out otherwise. continue except RuntimeError as e: if configuration['mpi'] is True: raise RuntimeError(str(e)) for k, v in processed.items(): insert(v, k.parent, [k]) return stree