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
