def _schedule_expressions(self, clusters):
"""Create an Iteartion/Expression tree given an iterable of
:class:`Cluster` objects."""
# Build the Iteration/Expression tree
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 or j0.dim in clusters.atomics[i]:
break
root = schedule[j1]
index += 1
needed = entries[index:]
# Build and insert the required Iterations
iters = [
Iteration([], j.dim, j.dim.limits, 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 List(body=processed)
def make_blocking(self, iet):
"""
Apply loop blocking to PARALLEL Iteration trees.
"""
# Make sure loop blocking will span as many Iterations as possible
iet = fold_blockable_tree(iet, self.blockinner)
mapper = {}
efuncs = []
block_dims = []
for tree in retrieve_iteration_tree(iet):
# Is the Iteration tree blockable ?
iterations = filter_iterations(tree, lambda i: i.is_Parallel and i.is_Affine)
if not self.blockinner:
iterations = iterations[:-1]
if len(iterations) <= 1:
continue
root = iterations[0]
if not self.blockalways:
# Heuristically bypass loop blocking if we think `tree`
# won't be computationally expensive. This will help with code
# size/readbility, JIT time, and auto-tuning time
if not (tree.root.is_Sequential or iet.is_Callable):
# E.g., not inside a time-stepping Iteration
continue
if any(i.dim.is_Sub and i.dim.local for i in tree):
# At least an outer Iteration is over a local SubDimension,
# which suggests the computational cost of this Iteration
# nest will be negligible w.r.t. the "core" Iteration nest
# (making use of non-local (Sub)Dimensions only)
continue
if not IsPerfectIteration().visit(root):
# Don't know how to block non-perfect nests
continue
# Apply hierarchical loop blocking to `tree`
level_0 = [] # Outermost level of blocking
level_i = [[] for i in range(1, self.nlevels)] # Inner levels of blocking
intra = [] # Within the smallest block
for i in iterations:
template = "%s%d_blk%s" % (i.dim.name, self.nblocked, '%d')
properties = (PARALLEL,) + ((AFFINE,) if i.is_Affine else ())
# Build Iteration across `level_0` blocks
d = BlockDimension(i.dim, name=template % 0)
level_0.append(Iteration([], d, d.symbolic_max, properties=properties))
# Build Iteration across all `level_i` blocks, `i` in (1, self.nlevels]
for n, li in enumerate(level_i, 1):
di = BlockDimension(d, name=template % n)
li.append(Iteration([], di, limits=(d, d+d.step-1, di.step),
properties=properties))
d = di
# Build Iteration within the smallest block
intra.append(i._rebuild([], limits=(d, d+d.step-1, 1), offsets=(0, 0)))
level_i = flatten(level_i)
# Track all constructed BlockDimensions
block_dims.extend(i.dim for i in level_0 + level_i)
# Construct the blocked tree
blocked = compose_nodes(level_0 + level_i + intra + [iterations[-1].nodes])
blocked = unfold_blocked_tree(blocked)
# Promote to a separate Callable
dynamic_parameters = flatten((l0.dim, l0.step) for l0 in level_0)
dynamic_parameters.extend([li.step for li in level_i])
efunc = make_efunc("bf%d" % self.nblocked, blocked, dynamic_parameters)
efuncs.append(efunc)
# Compute the iteration ranges
ranges = []
for i, l0 in zip(iterations, level_0):
maxb = i.symbolic_max - (i.symbolic_size % l0.step)
ranges.append(((i.symbolic_min, maxb, l0.step),
(maxb + 1, i.symbolic_max, i.symbolic_max - maxb)))
# Build Calls to the `efunc`
body = []
for p in product(*ranges):
dynamic_args_mapper = {}
for l0, (m, M, b) in zip(level_0, p):
dynamic_args_mapper[l0.dim] = (m, M)
dynamic_args_mapper[l0.step] = (b,)
for li in level_i:
if li.dim.root is l0.dim.root:
value = li.step if b is l0.step else b
dynamic_args_mapper[li.step] = (value,)
call = efunc.make_call(dynamic_args_mapper)
body.append(List(body=call))
mapper[root] = List(body=body)
# Next blockable nest, use different (unique) variable/function names
self.nblocked += 1
iet = Transformer(mapper).visit(iet)
# Force-unfold if some folded Iterations haven't been blocked in the end
iet = unfold_blocked_tree(iet)
return iet, {'dimensions': block_dims,
'efuncs': efuncs,
'args': [i.step for i in block_dims]}
def optimize_unfolded_tree(unfolded, root):
"""
Transform folded trees to reduce the memory footprint.
Examples
--------
Given:
.. code-block::
for i = 1 to N - 1 # Folded tree
for j = 1 to N - 1
tmp[i,j] = ...
for i = 2 to N - 2 # Root
for j = 2 to N - 2
... = ... tmp[i,j] ...
The temporary ``tmp`` has shape ``(N-1, N-1)``. However, as soon as the
iteration space is blocked, with blocks of shape ``(i_bs, j_bs)``, the
``tmp`` shape can be shrunk to ``(i_bs-1, j_bs-1)``. The resulting
iteration tree becomes:
.. code-block::
for i = 1 to i_bs + 1 # Folded tree
for j = 1 to j_bs + 1
i' = i + i_block - 2
j' = j + j_block - 2
tmp[i,j] = ... # use i' and j'
for i = i_block to i_block + i_bs # Root
for j = j_block to j_block + j_bs
i' = i - x_block
j' = j - j_block
... = ... tmp[i',j'] ...
"""
processed = []
for i, tree in enumerate(unfolded):
assert len(tree) == len(root)
# We can optimize the folded trees only iff:
# test0 := they compute temporary arrays, but not if they compute input data
# test1 := the outer Iterations have actually been blocked
exprs = FindNodes(Expression).visit(tree)
writes = [j.write for j in exprs if j.is_tensor]
test0 = not all(j.is_Array for j in writes)
test1 = any(not isinstance(j.limits[0], BlockDimension) for j in root)
if test0 or test1:
processed.append(compose_nodes(tree))
root = compose_nodes(root)
continue
# Shrink the iteration space
modified_tree = []
modified_root = []
modified_dims = {}
mapper = {}
for t, r in zip(tree, root):
udim0 = IncrDimension(t.dim, t.symbolic_min, 1, "%ss%d" % (t.index, i))
modified_tree.append(t._rebuild(limits=(0, t.limits[1] - t.limits[0], t.step),
uindices=t.uindices + (udim0,)))
mapper[t.dim] = udim0
udim1 = IncrDimension(t.dim, 0, 1, "%ss%d" % (t.index, i))
modified_root.append(r._rebuild(uindices=r.uindices + (udim1,)))
d = r.limits[0]
assert isinstance(d, BlockDimension)
modified_dims[d.root] = d
# Temporary arrays can now be moved onto the stack
for w in writes:
dims = tuple(modified_dims.get(d, d) for d in w.dimensions)
shape = tuple(d.symbolic_size for d in dims)
w.update(shape=shape, dimensions=dims, scope='stack')
# Substitute iteration variables within the folded trees
modified_tree = compose_nodes(modified_tree)
replaced = xreplace_indices([j.expr for j in exprs], mapper,
lambda i: i.function not in writes, True)
subs = [j._rebuild(expr=k) for j, k in zip(exprs, replaced)]
processed.append(Transformer(dict(zip(exprs, subs))).visit(modified_tree))
# Introduce the new iteration variables within /root/
modified_root = compose_nodes(modified_root)
exprs = FindNodes(Expression).visit(modified_root)
candidates = [as_symbol(j.output) for j in subs]
replaced = xreplace_indices([j.expr for j in exprs], mapper, candidates)
subs = [j._rebuild(expr=k) for j, k in zip(exprs, replaced)]
root = Transformer(dict(zip(exprs, subs))).visit(modified_root)
return processed + [root]
def relax_incr_dimensions(iet, **kwargs):
"""
Recast Iterations over IncrDimensions as ElementalFunctions; insert
ElementalCalls to iterate over the "main" and "remainder" regions induced
by the IncrDimensions.
"""
sregistry = kwargs['sregistry']
efuncs = []
mapper = {}
for tree in retrieve_iteration_tree(iet):
iterations = [i for i in tree if i.dim.is_Incr]
if not iterations:
continue
root = iterations[0]
if root in mapper:
continue
outer, inner = split(iterations, lambda i: not i.dim.parent.is_Incr)
# Compute the iteration ranges
ranges = []
for i in outer:
maxb = i.symbolic_max - (i.symbolic_size % i.dim.step)
ranges.append(((i.symbolic_min, maxb, i.dim.step),
(maxb + 1, i.symbolic_max, i.symbolic_max - maxb)))
# Remove any offsets
# E.g., `x = x_m + 2 to x_M - 2` --> `x = x_m to x_M`
outer = [i._rebuild(limits=(i.dim.root.symbolic_min, i.dim.root.symbolic_max,
i.step))
for i in outer]
# Create the ElementalFunction
name = sregistry.make_name(prefix="bf")
body = compose_nodes(outer)
dynamic_parameters = flatten((i.symbolic_bounds, i.step) for i in outer)
dynamic_parameters.extend([i.step for i in inner if not is_integer(i.step)])
efunc = make_efunc(name, body, dynamic_parameters)
efuncs.append(efunc)
# Create the ElementalCalls
calls = []
for p in product(*ranges):
dynamic_args_mapper = {}
for i, (m, M, b) in zip(outer, p):
dynamic_args_mapper[i.symbolic_min] = m
dynamic_args_mapper[i.symbolic_max] = M
dynamic_args_mapper[i.step] = b
for j in inner:
if j.dim.root is i.dim.root and not is_integer(j.step):
value = j.step if b is i.step else b
dynamic_args_mapper[j.step] = (value,)
calls.append(efunc.make_call(dynamic_args_mapper))
mapper[root] = List(body=calls)
iet = Transformer(mapper).visit(iet)
return iet, {'efuncs': efuncs}
def optimize_unfolded_tree(unfolded, root):
"""
Transform folded trees to reduce the memory footprint.
Examples
========
Given:
.. code-block::
for i = 1 to N - 1 # Folded tree
for j = 1 to N - 1
tmp[i,j] = ...
for i = 2 to N - 2 # Root
for j = 2 to N - 2
... = ... tmp[i,j] ...
The temporary ``tmp`` has shape ``(N-1, N-1)``. However, as soon as the
iteration space is blocked, with blocks of shape ``(i_bs, j_bs)``, the
``tmp`` shape can be shrunk to ``(i_bs-1, j_bs-1)``. The resulting
iteration tree becomes:
.. code-block::
for i = 1 to i_bs + 1 # Folded tree
for j = 1 to j_bs + 1
i' = i + i_block - 2
j' = j + j_block - 2
tmp[i,j] = ... # use i' and j'
for i = i_block to i_block + i_bs # Root
for j = j_block to j_block + j_bs
i' = i - x_block
j' = j - j_block
... = ... tmp[i',j'] ...
"""
processed = []
for i, tree in enumerate(unfolded):
assert len(tree) == len(root)
modified_tree = []
modified_root = []
mapper = {}
# "Shrink" the iteration space
for t1, t2 in zip(tree, root):
index = Symbol('%ss%d' % (t1.index, i))
mapper[t1.dim] = index
t1_uindex = (UnboundedIndex(index, t1.limits[0]), )
t2_uindex = (UnboundedIndex(index, -t1.limits[0]), )
limits = (0, t1.limits[1] - t1.limits[0], t1.incr_symbolic)
modified_tree.append(
t1._rebuild(limits=limits, uindices=t1.uindices + t1_uindex))
modified_root.append(t2._rebuild(uindices=t2.uindices + t2_uindex))
# Temporary arrays can now be moved onto the stack
exprs = FindNodes(Expression).visit(modified_tree[-1])
if all(not j.is_Remainder for j in modified_tree):
dimensions = tuple(j.limits[0] for j in modified_root)
for j in exprs:
if j.write.is_Array:
j_dimensions = dimensions + j.write.dimensions[
len(modified_root):]
j_shape = tuple(k.symbolic_size for k in j_dimensions)
j.write.update(shape=j_shape,
dimensions=j_dimensions,
onstack=True)
# Substitute iteration variables within the folded trees
modified_tree = compose_nodes(modified_tree)
replaced = xreplace_indices([j.expr for j in exprs],
mapper,
only_rhs=True)
subs = [j._rebuild(expr=k) for j, k in zip(exprs, replaced)]
processed.append(
Transformer(dict(zip(exprs, subs))).visit(modified_tree))
# Introduce the new iteration variables within /root/
modified_root = compose_nodes(modified_root)
exprs = FindNodes(Expression).visit(modified_root)
candidates = [as_symbol(j.output) for j in subs]
replaced = xreplace_indices([j.expr for j in exprs], mapper,
candidates)
subs = [j._rebuild(expr=k) for j, k in zip(exprs, replaced)]
root = Transformer(dict(zip(exprs, subs))).visit(modified_root)
return processed + [root]
def _minimize_remainders(self, iet):
"""
Reshape temporary tensors and adjust loop trip counts to prevent as many
compiler-generated remainder loops as possible.
"""
# The innermost dimension is the one that might get padded
p_dim = -1
mapper = {}
for tree in retrieve_iteration_tree(iet):
vector_iterations = [i for i in tree if i.is_Vectorizable]
if not vector_iterations or len(vector_iterations) > 1:
continue
root = vector_iterations[0]
# Padding
writes = [i.write for i in FindNodes(Expression).visit(root)
if i.write.is_Array]
padding = []
for i in writes:
try:
simd_items = self.platform.simd_items_per_reg(i.dtype)
except KeyError:
return iet, {}
padding.append(simd_items - i.shape[-1] % simd_items)
if len(set(padding)) == 1:
padding = padding[0]
for i in writes:
padded = (i._padding[p_dim][0], i._padding[p_dim][1] + padding)
i.update(padding=i._padding[:p_dim] + (padded,))
else:
# Padding must be uniform -- not the case, so giving up
continue
# Dynamic trip count adjustment
endpoint = root.symbolic_max
if not endpoint.is_Symbol:
continue
condition = []
externals = set(i.symbolic_shape[-1] for i in FindSymbols().visit(root)
if i.is_Tensor)
for i in root.uindices:
for j in externals:
condition.append(root.symbolic_max + padding < j)
condition = ' && '.join(ccode(i) for i in condition)
endpoint_padded = endpoint.func('_%s' % endpoint.name)
init = cgen.Initializer(
cgen.Value("const int", endpoint_padded),
cgen.Line('(%s) ? %s : %s' % (condition,
ccode(endpoint + padding),
endpoint))
)
# Update the Iteration bound
limits = list(root.limits)
limits[1] = endpoint_padded.func(endpoint_padded.name)
rebuilt = list(tree)
rebuilt[rebuilt.index(root)] = root._rebuild(limits=limits)
mapper[tree[0]] = List(header=init, body=compose_nodes(rebuilt))
processed = Transformer(mapper).visit(iet)
return processed, {}
def _loop_blocking(self, iet):
"""
Apply loop blocking to PARALLEL Iteration trees.
"""
blockinner = bool(self.params.get('blockinner'))
blockalways = bool(self.params.get('blockalways'))
# Make sure loop blocking will span as many Iterations as possible
iet = fold_blockable_tree(iet, blockinner)
mapper = {}
efuncs = []
block_dims = []
for tree in retrieve_iteration_tree(iet):
# Is the Iteration tree blockable ?
iterations = filter_iterations(tree, lambda i: i.is_Parallel)
if not blockinner:
iterations = iterations[:-1]
if len(iterations) <= 1:
continue
root = iterations[0]
if not blockalways:
# Heuristically bypass loop blocking if we think `tree`
# won't be computationally expensive. This will help with code
# size/redability, JIT time, and auto-tuning time
if not (tree.root.is_Sequential or iet.is_Callable):
# E.g., not inside a time-stepping Iteration
continue
if any(i.dim.is_Sub and i.dim.local for i in tree):
# At least an outer Iteration is over a local SubDimension,
# which suggests the computational cost of this Iteration
# nest will be negligible w.r.t. the "core" Iteration nest
# (making use of non-local (Sub)Dimensions only)
continue
if not IsPerfectIteration().visit(root):
# Don't know how to block non-perfect nests
continue
# Apply loop blocking to `tree`
interb = []
intrab = []
for i in iterations:
d = BlockDimension(i.dim, name="%s%d_blk" % (i.dim.name, len(mapper)))
block_dims.append(d)
# Build Iteration over blocks
properties = (PARALLEL,) + ((AFFINE,) if i.is_Affine else ())
interb.append(Iteration([], d, d.symbolic_max, properties=properties))
# Build Iteration within a block
intrab.append(i._rebuild([], limits=(d, d+d.step-1, 1), offsets=(0, 0)))
# Construct the blocked tree
blocked = compose_nodes(interb + intrab + [iterations[-1].nodes])
blocked = unfold_blocked_tree(blocked)
# Promote to a separate Callable
dynamic_parameters = flatten((bi.dim, bi.dim.symbolic_size) for bi in interb)
efunc = make_efunc("bf%d" % len(mapper), blocked, dynamic_parameters)
efuncs.append(efunc)
# Compute the iteration ranges
ranges = []
for i, bi in zip(iterations, interb):
maxb = i.symbolic_max - (i.symbolic_size % bi.dim.step)
ranges.append(((i.symbolic_min, maxb, bi.dim.step),
(maxb + 1, i.symbolic_max, i.symbolic_max - maxb)))
# Build Calls to the `efunc`
body = []
for p in product(*ranges):
dynamic_args_mapper = {}
for bi, (m, M, b) in zip(interb, p):
dynamic_args_mapper[bi.dim] = (m, M)
dynamic_args_mapper[bi.dim.step] = (b,)
call = efunc.make_call(dynamic_args_mapper)
body.append(List(body=call))
mapper[root] = List(body=body)
iet = Transformer(mapper).visit(iet)
return iet, {'dimensions': block_dims, 'efuncs': efuncs,
'args': [i.step for i in block_dims]}
def _loop_blocking(self, iet):
"""
Apply loop blocking to PARALLEL Iteration trees.
"""
blockinner = bool(self.params.get('blockinner'))
blockalways = bool(self.params.get('blockalways'))
# Make sure loop blocking will span as many Iterations as possible
iet = fold_blockable_tree(iet, blockinner)
mapper = {}
efuncs = []
block_dims = []
for tree in retrieve_iteration_tree(iet):
# Is the Iteration tree blockable ?
iterations = filter_iterations(tree, lambda i: i.is_Parallel)
if not blockinner:
iterations = iterations[:-1]
if len(iterations) <= 1:
continue
root = iterations[0]
if not (tree.root.is_Sequential
or iet.is_Callable) and not blockalways:
# Heuristic: avoid polluting the generated code with blocked
# nests (thus increasing JIT compilation time and affecting
# readability) if the blockable tree isn't embedded in a
# sequential loop (e.g., a timestepping loop)
continue
# Apply loop blocking to `tree`
interb = []
intrab = []
for i in iterations:
d = BlockDimension(i.dim,
name="%s%d_blk" % (i.dim.name, len(mapper)))
block_dims.append(d)
# Build Iteration over blocks
interb.append(
Iteration([], d, d.symbolic_max, properties=PARALLEL))
# Build Iteration within a block
intrab.append(
i._rebuild([],
limits=(d, d + d.step - 1, 1),
offsets=(0, 0)))
# Construct the blocked tree
blocked = compose_nodes(interb + intrab + [iterations[-1].nodes])
blocked = unfold_blocked_tree(blocked)
# Promote to a separate Callable
dynamic_parameters = flatten(
(bi.dim, bi.dim.symbolic_size) for bi in interb)
efunc = make_efunc("bf%d" % len(mapper), blocked,
dynamic_parameters)
efuncs.append(efunc)
# Compute the iteration ranges
ranges = []
for i, bi in zip(iterations, interb):
maxb = i.symbolic_max - (i.symbolic_size % bi.dim.step)
ranges.append(
((i.symbolic_min, maxb, bi.dim.step),
(maxb + 1, i.symbolic_max, i.symbolic_max - maxb)))
# Build Calls to the `efunc`
body = []
for p in product(*ranges):
dynamic_args_mapper = {}
for bi, (m, M, b) in zip(interb, p):
dynamic_args_mapper[bi.dim] = (m, M)
dynamic_args_mapper[bi.dim.step] = (b, )
call = efunc.make_call(dynamic_args_mapper)
body.append(List(body=call))
mapper[root] = List(body=body)
iet = Transformer(mapper).visit(iet)
return iet, {
'dimensions': block_dims,
'efuncs': efuncs,
'args': [i.step for i in block_dims]
}
def _loop_blocking(self, nodes, state):
"""
Apply loop blocking to :class:`Iteration` trees.
Blocking is applied to parallel iteration trees. Heuristically, innermost
dimensions are not blocked to maximize the trip count of the SIMD loops.
Different heuristics may be specified by passing the keywords ``blockshape``
and ``blockinner`` to the DLE. The former, a dictionary, is used to indicate
a specific block size for each blocked dimension. For example, for the
:class:`Iteration` tree: ::
for i
for j
for k
...
one may provide ``blockshape = {i: 4, j: 7}``, in which case the
two outer loops will blocked, and the resulting 2-dimensional block will
have size 4x7. The latter may be set to True to also block innermost parallel
:class:`Iteration` objects.
"""
exclude_innermost = not self.params.get('blockinner', False)
ignore_heuristic = self.params.get('blockalways', False)
# Make sure loop blocking will span as many Iterations as possible
fold = fold_blockable_tree(nodes, exclude_innermost)
mapper = {}
blocked = OrderedDict()
for tree in retrieve_iteration_tree(fold):
# Is the Iteration tree blockable ?
iterations = [i for i in tree if i.is_Parallel]
if exclude_innermost:
iterations = [i for i in iterations if not i.is_Vectorizable]
if len(iterations) <= 1:
continue
root = iterations[0]
if not IsPerfectIteration().visit(root):
# Illegal/unsupported
continue
if not tree[0].is_Sequential and not ignore_heuristic:
# Heuristic: avoid polluting the generated code with blocked
# nests (thus increasing JIT compilation time and affecting
# readability) if the blockable tree isn't embedded in a
# sequential loop (e.g., a timestepping loop)
continue
# Decorate intra-block iterations with an IterationProperty
TAG = tagger(len(mapper))
# Build all necessary Iteration objects, individually. These will
# subsequently be composed to implement loop blocking.
inter_blocks = []
intra_blocks = []
remainders = []
for i in iterations:
name = "%s%d_block" % (i.dim.name, len(mapper))
# Build Iteration over blocks
dim = blocked.setdefault(i, Dimension(name=name))
bsize = dim.symbolic_size
bstart = i.limits[0]
binnersize = i.dim.symbolic_extent + (i.offsets[1] -
i.offsets[0])
bfinish = i.dim.symbolic_end - (binnersize % bsize) - 1
inter_block = Iteration([],
dim, [bstart, bfinish, bsize],
offsets=i.offsets,
properties=PARALLEL)
inter_blocks.append(inter_block)
# Build Iteration within a block
limits = (dim, dim + bsize - 1, 1)
intra_block = i._rebuild([],
limits=limits,
offsets=(0, 0),
properties=i.properties +
(TAG, ELEMENTAL))
intra_blocks.append(intra_block)
# Build unitary-increment Iteration over the 'leftover' region.
# This will be used for remainder loops, executed when any
# dimension size is not a multiple of the block size.
remainder = i._rebuild(
[],
limits=[bfinish + 1, i.dim.symbolic_end, 1],
offsets=(i.offsets[1], i.offsets[1]))
remainders.append(remainder)
# Build blocked Iteration nest
blocked_tree = compose_nodes(inter_blocks + intra_blocks +
[iterations[-1].nodes])
# Build remainder Iterations
remainder_trees = []
for n in range(len(iterations)):
for c in combinations([i.dim for i in iterations], n + 1):
# First all inter-block Interations
nodes = [
b._rebuild(properties=b.properties + (REMAINDER, ))
for b, r in zip(inter_blocks, remainders)
if r.dim not in c
]
# Then intra-block or remainder, for each dim (in order)
properties = (REMAINDER, TAG, ELEMENTAL)
for b, r in zip(intra_blocks, remainders):
handle = r if b.dim in c else b
nodes.append(handle._rebuild(properties=properties))
nodes.extend([iterations[-1].nodes])
remainder_trees.append(compose_nodes(nodes))
# Will replace with blocked loop tree
mapper[root] = List(body=[blocked_tree] + remainder_trees)
rebuilt = Transformer(mapper).visit(fold)
# Finish unrolling any previously folded Iterations
processed = unfold_blocked_tree(rebuilt)
# All blocked dimensions
if not blocked:
return processed, {}
# Determine the block shape
blockshape = self.params.get('blockshape')
if not blockshape:
# Use trivial heuristic for a suitable blockshape
def heuristic(dim_size):
ths = 8 # FIXME: This really needs to be improved
return ths if dim_size > ths else 1
blockshape = {k: heuristic for k in blocked.keys()}
else:
try:
nitems, nrequired = len(blockshape), len(blocked)
blockshape = {k: v for k, v in zip(blocked, blockshape)}
if nitems > nrequired:
dle_warning("Provided 'blockshape' has more entries than "
"blocked loops; dropping entries ...")
if nitems < nrequired:
dle_warning("Provided 'blockshape' has fewer entries than "
"blocked loops; dropping dimensions ...")
except TypeError:
blockshape = {list(blocked)[0]: blockshape}
blockshape.update(
{k: None
for k in blocked.keys() if k not in blockshape})
# Track any additional arguments required to execute /state.nodes/
arguments = [
BlockingArg(v, k, blockshape[k]) for k, v in blocked.items()
]
return processed, {'arguments': arguments, 'flags': 'blocking'}
def _loop_blocking(self, nodes, state):
"""Apply loop blocking to PARALLEL Iteration trees."""
exclude_innermost = not self.params.get('blockinner', False)
ignore_heuristic = self.params.get('blockalways', False)
# Make sure loop blocking will span as many Iterations as possible
fold = fold_blockable_tree(nodes, exclude_innermost)
mapper = {}
blocked = OrderedDict()
for tree in retrieve_iteration_tree(fold):
# Is the Iteration tree blockable ?
iterations = [i for i in tree if i.is_Parallel]
if exclude_innermost:
iterations = [i for i in iterations if not i.is_Vectorizable]
if len(iterations) <= 1:
continue
root = iterations[0]
if not IsPerfectIteration().visit(root):
# Illegal/unsupported
continue
if not tree.root.is_Sequential and not ignore_heuristic:
# Heuristic: avoid polluting the generated code with blocked
# nests (thus increasing JIT compilation time and affecting
# readability) if the blockable tree isn't embedded in a
# sequential loop (e.g., a timestepping loop)
continue
# Decorate intra-block iterations with an IterationProperty
TAG = tagger(len(mapper))
# Build all necessary Iteration objects, individually. These will
# subsequently be composed to implement loop blocking.
inter_blocks = []
intra_blocks = []
remainders = []
for i in iterations:
# Build Iteration over blocks
name = "%s%d_block" % (i.dim.name, len(mapper))
dim = blocked.setdefault(i, BlockDimension(i.dim, name=name))
binnersize = i.symbolic_size + (i.offsets[1] - i.offsets[0])
bmax = i.dim.symbolic_max - (binnersize % dim.step)
inter_block = Iteration([], dim, bmax, offsets=i.offsets,
properties=PARALLEL)
inter_blocks.append(inter_block)
# Build Iteration within a block
limits = (dim, dim + dim.step - 1, 1)
intra_block = i._rebuild([], limits=limits, offsets=(0, 0),
properties=i.properties + (TAG, ELEMENTAL))
intra_blocks.append(intra_block)
# Build unitary-increment Iteration over the 'leftover' region.
# This will be used for remainder loops, executed when any
# dimension size is not a multiple of the block size.
remainder = i._rebuild([], limits=[bmax + 1, i.dim.symbolic_max, 1],
offsets=(i.offsets[1], i.offsets[1]))
remainders.append(remainder)
# Build blocked Iteration nest
blocked_tree = compose_nodes(inter_blocks + intra_blocks +
[iterations[-1].nodes])
# Build remainder Iterations
remainder_trees = []
for n in range(len(iterations)):
for c in combinations([i.dim for i in iterations], n + 1):
# First all inter-block Interations
nodes = [b._rebuild(properties=b.properties + (REMAINDER,))
for b, r in zip(inter_blocks, remainders)
if r.dim not in c]
# Then intra-block or remainder, for each dim (in order)
properties = (REMAINDER, TAG, ELEMENTAL)
for b, r in zip(intra_blocks, remainders):
handle = r if b.dim in c else b
nodes.append(handle._rebuild(properties=properties))
nodes.extend([iterations[-1].nodes])
remainder_trees.append(compose_nodes(nodes))
# Will replace with blocked loop tree
mapper[root] = List(body=[blocked_tree] + remainder_trees)
rebuilt = Transformer(mapper).visit(fold)
# Finish unrolling any previously folded Iterations
processed = unfold_blocked_tree(rebuilt)
return processed, {'dimensions': list(blocked.values())}
def _minimize_remainders(self, iet):
"""
Reshape temporary tensors and adjust loop trip counts to prevent as many
compiler-generated remainder loops as possible.
"""
# The innermost dimension is the one that might get padded
p_dim = -1
mapper = {}
for tree in retrieve_iteration_tree(iet):
vector_iterations = [i for i in tree if i.is_Vectorizable]
if not vector_iterations or len(vector_iterations) > 1:
continue
root = vector_iterations[0]
# Padding
writes = [i.write for i in FindNodes(Expression).visit(root)
if i.write.is_Array]
padding = []
for i in writes:
try:
simd_items = self.platform.simd_items_per_reg(i.dtype)
except KeyError:
return iet, {}
padding.append(simd_items - i.shape[-1] % simd_items)
if len(set(padding)) == 1:
padding = padding[0]
for i in writes:
padded = (i._padding[p_dim][0], i._padding[p_dim][1] + padding)
i.update(padding=i._padding[:p_dim] + (padded,))
else:
# Padding must be uniform -- not the case, so giving up
continue
# Dynamic trip count adjustment
endpoint = root.symbolic_max
if not endpoint.is_Symbol:
continue
condition = []
externals = set(i.symbolic_shape[-1] for i in FindSymbols().visit(root)
if i.is_Tensor)
for i in root.uindices:
for j in externals:
condition.append(root.symbolic_max + padding < j)
condition = ' && '.join(ccode(i) for i in condition)
endpoint_padded = endpoint.func('_%s' % endpoint.name)
init = cgen.Initializer(
cgen.Value("const int", endpoint_padded),
cgen.Line('(%s) ? %s : %s' % (condition,
ccode(endpoint + padding),
endpoint))
)
# Update the Iteration bound
limits = list(root.limits)
limits[1] = endpoint_padded.func(endpoint_padded.name)
rebuilt = list(tree)
rebuilt[rebuilt.index(root)] = root._rebuild(limits=limits)
mapper[tree[0]] = List(header=init, body=compose_nodes(rebuilt))
processed = Transformer(mapper).visit(iet)
return processed, {}
def _loop_blocking(self, iet):
"""
Apply loop blocking to PARALLEL Iteration trees.
"""
blockinner = bool(self.params.get('blockinner'))
blockalways = bool(self.params.get('blockalways'))
# Make sure loop blocking will span as many Iterations as possible
iet = fold_blockable_tree(iet, blockinner)
mapper = {}
efuncs = []
block_dims = []
for tree in retrieve_iteration_tree(iet):
# Is the Iteration tree blockable ?
iterations = filter_iterations(tree, lambda i: i.is_Parallel)
if not blockinner:
iterations = iterations[:-1]
if len(iterations) <= 1:
continue
root = iterations[0]
if not (tree.root.is_Sequential or iet.is_Callable) and not blockalways:
# Heuristic: avoid polluting the generated code with blocked
# nests (thus increasing JIT compilation time and affecting
# readability) if the blockable tree isn't embedded in a
# sequential loop (e.g., a timestepping loop)
continue
# Apply loop blocking to `tree`
interb = []
intrab = []
for i in iterations:
d = BlockDimension(i.dim, name="%s%d_blk" % (i.dim.name, len(mapper)))
block_dims.append(d)
# Build Iteration over blocks
properties = (PARALLEL,) + ((AFFINE,) if i.is_Affine else ())
interb.append(Iteration([], d, d.symbolic_max, properties=properties))
# Build Iteration within a block
intrab.append(i._rebuild([], limits=(d, d+d.step-1, 1), offsets=(0, 0)))
# Construct the blocked tree
blocked = compose_nodes(interb + intrab + [iterations[-1].nodes])
blocked = unfold_blocked_tree(blocked)
# Promote to a separate Callable
dynamic_parameters = flatten((bi.dim, bi.dim.symbolic_size) for bi in interb)
efunc = make_efunc("bf%d" % len(mapper), blocked, dynamic_parameters)
efuncs.append(efunc)
# Compute the iteration ranges
ranges = []
for i, bi in zip(iterations, interb):
maxb = i.symbolic_max - (i.symbolic_size % bi.dim.step)
ranges.append(((i.symbolic_min, maxb, bi.dim.step),
(maxb + 1, i.symbolic_max, i.symbolic_max - maxb)))
# Build Calls to the `efunc`
body = []
for p in product(*ranges):
dynamic_args_mapper = {}
for bi, (m, M, b) in zip(interb, p):
dynamic_args_mapper[bi.dim] = (m, M)
dynamic_args_mapper[bi.dim.step] = (b,)
call = efunc.make_call(dynamic_args_mapper)
body.append(List(body=call))
mapper[root] = List(body=body)
iet = Transformer(mapper).visit(iet)
return iet, {'dimensions': block_dims, 'efuncs': efuncs,
'args': [i.step for i in block_dims]}
def _minimize_remainders(self, nodes, state):
"""
Reshape temporary tensors and adjust loop trip counts to prevent as many
compiler-generated remainder loops as possible.
"""
mapper = {}
for tree in retrieve_iteration_tree(nodes):
vector_iterations = [i for i in tree if i.is_Vectorizable]
if not vector_iterations or len(vector_iterations) > 1:
continue
root = vector_iterations[0]
if root.tag is None:
continue
# Padding
writes = [
i for i in FindSymbols('symbolics-writes').visit(root)
if i.is_Array
]
padding = []
for i in writes:
try:
simd_items = get_simd_items(i.dtype)
except KeyError:
# Fallback to 16 (maximum expectable padding, for AVX512 registers)
simd_items = simdinfo['avx512f'] / np.dtype(
i.dtype).itemsize
padding.append(simd_items - i.shape[-1] % simd_items)
if len(set(padding)) == 1:
padding = padding[0]
for i in writes:
i.update(shape=i.shape[:-1] + (i.shape[-1] + padding, ))
else:
# Padding must be uniform -- not the case, so giving up
continue
# Dynamic trip count adjustment
endpoint = root.end_symbolic
if not endpoint.is_Symbol:
continue
condition = []
externals = set(i.symbolic_shape[-1]
for i in FindSymbols().visit(root))
for i in root.uindices:
for j in externals:
condition.append(root.end_symbolic + padding < j)
condition = ' || '.join(ccode(i) for i in condition)
endpoint_padded = endpoint.func(name='_%s' % endpoint.name)
init = cgen.Initializer(
cgen.Value("const int", endpoint_padded),
cgen.Line('(%s) ? %s : %s' %
(condition, ccode(endpoint + padding), endpoint)))
# Update the Iteration bound
limits = list(root.limits)
limits[1] = endpoint_padded.func(endpoint_padded.name)
rebuilt = list(tree)
rebuilt[rebuilt.index(root)] = root._rebuild(limits=limits)
mapper[tree[0]] = List(header=init, body=compose_nodes(rebuilt))
processed = Transformer(mapper).visit(nodes)
return processed, {}
def _loop_blocking(self, iet):
"""
Apply loop blocking to PARALLEL Iteration trees.
"""
blockinner = bool(self.params.get('blockinner'))
blockalways = bool(self.params.get('blockalways'))
noinline = self._compiler_decoration('noinline', cgen.Comment('noinline?'))
# Make sure loop blocking will span as many Iterations as possible
iet = fold_blockable_tree(iet, blockinner)
mapper = {}
efuncs = OrderedDict()
block_dims = []
for tree in retrieve_iteration_tree(iet):
# Is the Iteration tree blockable ?
candidates = [i for i in tree if i.is_Parallel]
if blockinner:
iterations = candidates
else:
iterations = [i for i in candidates if not i.is_Vectorizable]
if len(iterations) <= 1:
continue
root = iterations[0]
if not IsPerfectIteration().visit(root):
# Illegal/unsupported
continue
if not tree.root.is_Sequential and not blockalways:
# Heuristic: avoid polluting the generated code with blocked
# nests (thus increasing JIT compilation time and affecting
# readability) if the blockable tree isn't embedded in a
# sequential loop (e.g., a timestepping loop)
continue
# Apply loop blocking to `tree`
interb = []
intrab = []
for i in iterations:
d = BlockDimension(i.dim, name="%s%d_block" % (i.dim.name, len(mapper)))
# Build Iteration over blocks
interb.append(Iteration([], d, d.symbolic_max, offsets=i.offsets,
properties=PARALLEL))
# Build Iteration within a block
intrab.append(i._rebuild([], limits=(d, d+d.step-1, 1), offsets=(0, 0)))
# Record that a new BlockDimension has been introduced
block_dims.append(d)
# Construct the blocked tree
blocked = compose_nodes(interb + intrab + [iterations[-1].nodes])
blocked = unfold_blocked_tree(blocked)
# Promote to a separate Callable
dynamic_parameters = flatten((bi.dim, bi.dim.symbolic_size) for bi in interb)
efunc0 = make_efunc("bf%d" % len(mapper), blocked, dynamic_parameters)
# Compute the iteration ranges
ranges = []
for i, bi in zip(iterations, interb):
maxb = i.symbolic_max - (i.symbolic_size % bi.dim.step)
ranges.append(((i.symbolic_min, maxb, bi.dim.step),
(maxb + 1, i.symbolic_max, i.symbolic_max - maxb)))
# Build Calls to the `efunc`
body = []
for p in product(*ranges):
dynamic_args_mapper = {}
for bi, (m, M, b) in zip(interb, p):
dynamic_args_mapper[bi.dim] = (m, M)
dynamic_args_mapper[bi.dim.step] = (b,)
call = efunc0.make_call(dynamic_args_mapper)
body.append(List(header=noinline, body=call))
# Build indirect Call to the `efunc0` Calls
dynamic_parameters = [i.dim.root for i in candidates]
dynamic_parameters.extend([bi.dim.step for bi in interb])
efunc1 = make_efunc("f%d" % len(mapper), body, dynamic_parameters)
# Track everything to ultimately transform the input `iet`
mapper[root] = efunc1.make_call()
efuncs[efunc1] = None
efuncs[efunc0] = [efunc1.name]
iet = Transformer(mapper).visit(iet)
return iet, {'dimensions': block_dims, 'efuncs': efuncs}
def optimize_unfolded_tree(unfolded, root):
"""
Transform folded trees to reduce the memory footprint.
Examples
--------
Given:
.. code-block::
for i = 1 to N - 1 # Folded tree
for j = 1 to N - 1
tmp[i,j] = ...
for i = 2 to N - 2 # Root
for j = 2 to N - 2
... = ... tmp[i,j] ...
The temporary ``tmp`` has shape ``(N-1, N-1)``. However, as soon as the
iteration space is blocked, with blocks of shape ``(i_bs, j_bs)``, the
``tmp`` shape can be shrunk to ``(i_bs-1, j_bs-1)``. The resulting
iteration tree becomes:
.. code-block::
for i = 1 to i_bs + 1 # Folded tree
for j = 1 to j_bs + 1
i' = i + i_block - 2
j' = j + j_block - 2
tmp[i,j] = ... # use i' and j'
for i = i_block to i_block + i_bs # Root
for j = j_block to j_block + j_bs
i' = i - x_block
j' = j - j_block
... = ... tmp[i',j'] ...
"""
processed = []
for i, tree in enumerate(unfolded):
assert len(tree) == len(root)
# We can optimize the folded trees only if they compute temporary
# arrays, but not if they compute input data
exprs = FindNodes(Expression).visit(tree[-1])
writes = [j.write for j in exprs if j.is_tensor]
if not all(j.is_Array for j in writes):
processed.append(compose_nodes(tree))
root = compose_nodes(root)
continue
modified_tree = []
modified_root = []
mapper = {}
# "Shrink" the iteration space
for t1, t2 in zip(tree, root):
t1_udim = IncrDimension(t1.dim, t1.symbolic_min, 1, "%ss%d" % (t1.index, i))
limits = (0, t1.limits[1] - t1.limits[0], t1.step)
modified_tree.append(t1._rebuild(limits=limits,
uindices=t1.uindices + (t1_udim,)))
t2_udim = IncrDimension(t1.dim, 0, 1, "%ss%d" % (t1.index, i))
modified_root.append(t2._rebuild(uindices=t2.uindices + (t2_udim,)))
mapper[t1.dim] = t1_udim
# Temporary arrays can now be moved onto the stack
dimensions = tuple(j.limits[0] for j in modified_root)
for j in writes:
if j.is_Array:
j_dimensions = dimensions + j.dimensions[len(modified_root):]
j_shape = tuple(k.symbolic_size for k in j_dimensions)
j.update(shape=j_shape, dimensions=j_dimensions, scope='stack')
# Substitute iteration variables within the folded trees
modified_tree = compose_nodes(modified_tree)
replaced = xreplace_indices([j.expr for j in exprs], mapper, only_rhs=True)
subs = [j._rebuild(expr=k) for j, k in zip(exprs, replaced)]
processed.append(Transformer(dict(zip(exprs, subs))).visit(modified_tree))
# Introduce the new iteration variables within /root/
modified_root = compose_nodes(modified_root)
exprs = FindNodes(Expression).visit(modified_root)
candidates = [as_symbol(j.output) for j in subs]
replaced = xreplace_indices([j.expr for j in exprs], mapper, candidates)
subs = [j._rebuild(expr=k) for j, k in zip(exprs, replaced)]
root = Transformer(dict(zip(exprs, subs))).visit(modified_root)
return processed + [root]
def iet_make(clusters, dtype):
"""
Create an Iteration/Expression tree (IET) given an iterable of :class:`Cluster`s.
:param clusters: The iterable :class:`Cluster`s for which the IET is built.
:param dtype: The data type of the scalar expressions.
"""
processed = []
schedule = OrderedDict()
for cluster in clusters:
if not cluster.ispace.empty:
root = None
intervals = cluster.ispace.intervals
# Can I reuse any of the previously scheduled Iterations ?
index = 0
for i0, i1 in zip(intervals, list(schedule)):
if i0 != i1 or i0.dim in clusters.atomics[cluster]:
break
root = schedule[i1]
index += 1
needed = intervals[index:]
# Build Iterations, including any necessary unbounded index
iters = []
for i in needed:
uindices = []
for j, offs in cluster.ispace.sub_iterators.get(i.dim, []):
for n, o in enumerate(filter_ordered(offs)):
name = "%s%d" % (j.name, n)
vname = Scalar(name=name, dtype=np.int32)
value = (i.dim + o) % j.modulo
uindices.append(UnboundedIndex(vname, value, value, j, j + o))
iters.append(Iteration([], i.dim, i.dim.limits, offsets=i.limits,
uindices=uindices))
# Build Expressions
exprs = [Expression(v, np.int32 if cluster.trace.is_index(k) else dtype)
for k, v in cluster.trace.items()]
# Compose Iterations and Expressions
body, tree = compose_nodes(iters + [exprs], retrieve=True)
# Update the current scheduling
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([Expression(e, dtype) for e in cluster.exprs])
return List(body=processed)
