def iet_make(stree):
"""Create an IET from a ScheduleTree."""
nsections = 0
queues = OrderedDict()
for i in stree.visit():
if i == stree:
# We hit this handle at the very end of the visit
return List(body=queues.pop(i))
elif i.is_Exprs:
exprs = [Increment(e) if e.is_Increment else Expression(e) for e in i.exprs]
body = ExpressionBundle(i.ispace, i.ops, i.traffic, body=exprs)
elif i.is_Conditional:
body = Conditional(i.guard, queues.pop(i))
elif i.is_Iteration:
# Order to ensure deterministic code generation
uindices = sorted(i.sub_iterators, key=lambda d: d.name)
# Generate Iteration
body = Iteration(queues.pop(i), i.dim, i.limits, offsets=i.offsets,
direction=i.direction, properties=i.properties,
uindices=uindices)
elif i.is_Section:
body = Section('section%d' % nsections, body=queues.pop(i))
nsections += 1
elif i.is_Halo:
body = HaloSpot(i.halo_scheme, body=queues.pop(i))
queues.setdefault(i.parent, []).append(body)
assert False
def iet_build(stree):
"""
Construct an Iteration/Expression tree(IET) from a ScheduleTree.
"""
nsections = 0
queues = OrderedDict()
for i in stree.visit():
if i == stree:
# We hit this handle at the very end of the visit
return List(body=queues.pop(i))
elif i.is_Exprs:
exprs = [Increment(e) if e.is_Increment else Expression(e) for e in i.exprs]
body = ExpressionBundle(i.ispace, i.ops, i.traffic, body=exprs)
elif i.is_Conditional:
body = Conditional(i.guard, queues.pop(i))
elif i.is_Iteration:
body = Iteration(queues.pop(i), i.dim, i.limits, direction=i.direction,
properties=i.properties, uindices=i.sub_iterators)
elif i.is_Section:
body = Section('section%d' % nsections, body=queues.pop(i))
nsections += 1
elif i.is_Halo:
body = HaloSpot(i.halo_scheme, body=queues.pop(i))
queues.setdefault(i.parent, []).append(body)
assert False
def track_subsections(iet, **kwargs):
"""
Add custom Sections to the `profiler`. Custom Sections include:
* MPI Calls (e.g., HaloUpdateCall and HaloUpdateWait)
* Busy-waiting on While(lock) (e.g., from host-device orchestration)
"""
profiler = kwargs['profiler']
sregistry = kwargs['sregistry']
name_mapper = {
HaloUpdateCall: 'haloupdate',
HaloWaitCall: 'halowait',
RemainderCall: 'remainder',
HaloUpdateList: 'haloupdate',
HaloWaitList: 'halowait',
BusyWait: 'busywait'
}
mapper = {}
for NodeType in [MPIList, MPICall, BusyWait]:
for k, v in MapNodes(Section, NodeType).visit(iet).items():
for i in v:
if i in mapper or not any(
issubclass(i.__class__, n)
for n in profiler.trackable_subsections):
continue
name = sregistry.make_name(prefix=name_mapper[i.__class__])
mapper[i] = Section(name, body=i, is_subsection=True)
profiler.track_subsection(k.name, name)
iet = Transformer(mapper).visit(iet)
return iet, {}
def iet_make(stree):
"""
Create an Iteration/Expression tree (IET) from a :class:`ScheduleTree`.
"""
nsections = 0
queues = OrderedDict()
for i in stree.visit():
if i == stree:
# We hit this handle at the very end of the visit
return List(body=queues.pop(i))
elif i.is_Exprs:
exprs = [Expression(e) for e in i.exprs]
body = [ExpressionBundle(i.shape, i.ops, i.traffic, body=exprs)]
elif i.is_Conditional:
body = [Conditional(i.guard, queues.pop(i))]
elif i.is_Iteration:
# Order to ensure deterministic code generation
uindices = sorted(i.sub_iterators, key=lambda d: d.name)
# Generate Iteration
body = [Iteration(queues.pop(i), i.dim, i.dim.limits, offsets=i.limits,
direction=i.direction, uindices=uindices)]
elif i.is_Section:
body = [Section('section%d' % nsections, body=queues.pop(i))]
nsections += 1
elif i.is_Halo:
body = [HaloSpot(i.halo_scheme, body=queues.pop(i))]
queues.setdefault(i.parent, []).extend(body)
assert False
def iet_make(stree):
"""
Create an Iteration/Expression tree (IET) from a :class:`ScheduleTree`.
"""
nsections = 0
queues = OrderedDict()
for i in stree.visit():
if i == stree:
# We hit this handle at the very end of the visit
return List(body=queues.pop(i))
elif i.is_Exprs:
exprs = [Expression(e) for e in i.exprs]
body = [ExpressionBundle(i.shape, i.ops, i.traffic, body=exprs)]
elif i.is_Conditional:
body = [Conditional(i.guard, queues.pop(i))]
elif i.is_Iteration:
# Generate `uindices`
uindices = []
for d, offs in i.sub_iterators:
modulo = len(offs)
for n, o in enumerate(filter_ordered(offs)):
value = (i.dim + o) % modulo
symbol = Scalar(name="%s%d" % (d.name, n), dtype=np.int32)
uindices.append(
UnboundedIndex(symbol, value, value, d, d + o))
# Generate Iteration
body = [
Iteration(queues.pop(i),
i.dim,
i.dim.limits,
offsets=i.limits,
direction=i.direction,
uindices=uindices)
]
elif i.is_Section:
body = [Section('section%d' % nsections, body=queues.pop(i))]
nsections += 1
queues.setdefault(i.parent, []).extend(body)
assert False