def _schedule_expressions(self, clusters):
"""Wrap :class:`Expression` objects, already grouped in :class:`Cluster`
objects, within nested :class:`Iteration` objects (representing loops),
according to dimensions and stencils."""
# Topologically sort Iterations
ordering = partial_order([i.stencil.dimensions for i in clusters])
for i, d in enumerate(list(ordering)):
if d.is_Buffered:
ordering.insert(i, d.parent)
# Build the Iteration/Expression tree
processed = []
schedule = OrderedDict()
atomics = ()
for i in clusters:
# Build the Expression objects to be inserted within an Iteration tree
expressions = [Expression(v, np.int32 if i.trace.is_index(k) else self.dtype)
for k, v in i.trace.items()]
if not i.stencil.empty:
root = None
entries = i.stencil.entries
# Reorder based on the globally-established loop ordering
entries = sorted(entries, key=lambda i: ordering.index(i.dim))
# Can I reuse any of the previously scheduled Iterations ?
index = 0
for j0, j1 in zip(entries, list(schedule)):
if j0 != j1 or j0.dim in atomics:
break
root = schedule[j1]
index += 1
needed = entries[index:]
# Build and insert the required Iterations
iters = [Iteration([], j.dim, j.dim.size, offsets=j.ofs) for j in needed]
body, tree = compose_nodes(iters + [expressions], retrieve=True)
scheduling = OrderedDict(zip(needed, tree))
if root is None:
processed.append(body)
schedule = scheduling
else:
nodes = list(root.nodes) + [body]
mapper = {root: root._rebuild(nodes, **root.args_frozen)}
transformer = Transformer(mapper)
processed = list(transformer.visit(processed))
schedule = OrderedDict(list(schedule.items())[:index] +
list(scheduling.items()))
for k, v in list(schedule.items()):
schedule[k] = transformer.rebuilt.get(v, v)
else:
# No Iterations are needed
processed.extend(expressions)
# Track dimensions that cannot be fused at next stage
atomics = i.atomics
return List(body=processed)
def _schedule_expressions(self, clusters, ordering):
"""Wrap :class:`Expression` objects, already grouped in :class:`Cluster`
objects, within nested :class:`Iteration` objects (representing loops),
according to dimensions and stencils."""
processed = []
schedule = OrderedDict()
for i in clusters:
# Build the Expression objects to be inserted within an Iteration tree
expressions = [
Expression(v, np.int32 if i.trace.is_index(k) else self.dtype)
for k, v in i.trace.items()
]
if not i.stencil.empty:
root = None
entries = i.stencil.entries
# Can I reuse any of the previously scheduled Iterations ?
index = 0
for j0, j1 in zip(entries, list(schedule)):
if j0 != j1:
break
root = schedule[j1]
index += 1
needed = entries[index:]
# Build and insert the required Iterations
iters = [
Iteration([], j.dim, j.dim.size, offsets=j.ofs)
for j in needed
]
body, tree = compose_nodes(iters + [expressions],
retrieve=True)
scheduling = OrderedDict(zip(needed, tree))
if root is None:
processed.append(body)
schedule = scheduling
else:
nodes = list(root.nodes) + [body]
mapper = {root: root._rebuild(nodes, **root.args_frozen)}
transformer = Transformer(mapper)
processed = list(transformer.visit(processed))
schedule = OrderedDict(
list(schedule.items())[:index] +
list(scheduling.items()))
for k, v in list(schedule.items()):
schedule[k] = transformer.rebuilt.get(v, v)
else:
# No Iterations are needed
processed.extend(expressions)
return processed
def test_transformer_replace(exprs, block1, block2, block3):
"""Basic transformer test that replaces an expression"""
line1 = '// Replaced expression'
replacer = Block(c.Line(line1))
transformer = Transformer({exprs[0]: replacer})
for block in [block1, block2, block3]:
newblock = transformer.visit(block)
newcode = str(newblock.ccode)
oldnumlines = len(str(block.ccode).split('\n'))
newnumlines = len(newcode.split('\n'))
assert newnumlines >= oldnumlines
assert line1 in newcode
assert "a[i0] = a[i0] + b[i0] + 5.0F;" not in newcode
def test_transformer_wrap(exprs, block1, block2, block3):
"""Basic transformer test that wraps an expression in comments"""
line1 = '// This is the opening comment'
line2 = '// This is the closing comment'
wrapper = lambda n: Block(c.Line(line1), n, c.Line(line2))
transformer = Transformer({exprs[0]: wrapper(exprs[0])})
for block in [block1, block2, block3]:
newblock = transformer.visit(block)
newcode = str(newblock.ccode)
oldnumlines = len(str(block.ccode).split('\n'))
newnumlines = len(newcode.split('\n'))
assert newnumlines >= oldnumlines + 2
assert line1 in newcode
assert line2 in newcode
assert "a[i] = a[i] + b[i] + 5.0F;" in newcode
def test_transformer_add_replace(exprs, block2, block3):
"""Basic transformer test that adds one expression and replaces another"""
line1 = '// Replaced expression'
line2 = '// Adding a simple line'
replacer = Block(c.Line(line1))
adder = lambda n: Block(c.Line(line2), n)
transformer = Transformer({exprs[0]: replacer, exprs[1]: adder(exprs[1])})
for block in [block2, block3]:
newblock = transformer.visit(block)
newcode = str(newblock.ccode)
oldnumlines = len(str(block.ccode).split('\n'))
newnumlines = len(newcode.split('\n'))
assert newnumlines >= oldnumlines + 1
assert line1 in newcode
assert line2 in newcode
assert "a[i0] = a[i0] + b[i0] + 5.0F;" not in newcode