def apply(self, graph: SDFGState, sdfg: SDFG):
map_entry = self.map_entry
# Avoiding import loops
from dace.transformation.dataflow.strip_mining import StripMining
from dace.transformation.dataflow.local_storage import InLocalStorage, OutLocalStorage, LocalStorage
rangeexpr = str(map_entry.map.range.num_elements())
stripmine_subgraph = {
StripMining.map_entry: self.subgraph[MPITransformMap.map_entry]
}
sdfg_id = sdfg.sdfg_id
stripmine = StripMining()
stripmine.setup_match(sdfg, sdfg_id, self.state_id, stripmine_subgraph,
self.expr_index)
stripmine.dim_idx = -1
stripmine.new_dim_prefix = "mpi"
stripmine.tile_size = "(" + rangeexpr + "/__dace_comm_size)"
stripmine.divides_evenly = True
stripmine.apply(graph, sdfg)
# Find all in-edges that lead to the map entry
outer_map = None
edges = [
e for e in graph.in_edges(map_entry)
if isinstance(e.src, nodes.EntryNode)
]
outer_map = edges[0].src
# Add MPI schedule attribute to outer map
outer_map.map._schedule = dtypes.ScheduleType.MPI
# Now create a transient for each array
for e in edges:
if e.data.is_empty():
continue
in_local_storage_subgraph = {
LocalStorage.node_a: graph.node_id(outer_map),
LocalStorage.node_b: self.subgraph[MPITransformMap.map_entry]
}
sdfg_id = sdfg.sdfg_id
in_local_storage = InLocalStorage()
in_local_storage.setup_match(sdfg, sdfg_id, self.state_id,
in_local_storage_subgraph,
self.expr_index)
in_local_storage.array = e.data.data
in_local_storage.apply(graph, sdfg)
# Transform OutLocalStorage for each output of the MPI map
in_map_exit = graph.exit_node(map_entry)
out_map_exit = graph.exit_node(outer_map)
for e in graph.out_edges(out_map_exit):
if e.data.is_empty():
continue
name = e.data.data
outlocalstorage_subgraph = {
LocalStorage.node_a: graph.node_id(in_map_exit),
LocalStorage.node_b: graph.node_id(out_map_exit)
}
sdfg_id = sdfg.sdfg_id
outlocalstorage = OutLocalStorage()
outlocalstorage.setup_match(sdfg, sdfg_id, self.state_id,
outlocalstorage_subgraph,
self.expr_index)
outlocalstorage.array = name
outlocalstorage.apply(graph, sdfg)
def apply(self, sdfg):
graph = sdfg.node(self.state_id)
subgraph = self.subgraph_view(sdfg)
map_entries = helpers.get_outermost_scope_maps(sdfg, graph, subgraph)
result = StencilTiling.topology(sdfg, graph, map_entries)
(children_dict, parent_dict, sink_maps) = result
# next up, calculate inferred ranges for each map
# for each map entry, this contains a tuple of dicts:
# each of those maps from data_name of the array to
# inferred outer ranges. An inferred outer range is created
# by taking the union of ranges of inner subsets corresponding
# to that data and substituting this subset by the min / max of the
# parametrized map boundaries
# finally, from these outer ranges we can easily calculate
# strides and tile sizes required for every map
inferred_ranges = defaultdict(dict)
# create array of reverse topologically sorted map entries
# to iterate over
topo_reversed = []
queue = set(sink_maps.copy())
while len(queue) > 0:
element = next(e for e in queue if not children_dict[e] - set(topo_reversed))
topo_reversed.append(element)
queue.remove(element)
for parent in parent_dict[element]:
queue.add(parent)
# main loop
# first get coverage dicts for each map entry
# for each map, contains a tuple of two dicts
# each of those two maps from data name to outer range
coverage = {}
for map_entry in map_entries:
coverage[map_entry] = StencilTiling.coverage_dicts(sdfg, graph, map_entry, outer_range=True)
# we have a mapping from data name to outer range
# however we want a mapping from map parameters to outer ranges
# for this we need to find out how all array dimensions map to
# outer ranges
variable_mapping = defaultdict(list)
for map_entry in topo_reversed:
map = map_entry.map
# first find out variable mapping
for e in itertools.chain(graph.out_edges(map_entry), graph.in_edges(graph.exit_node(map_entry))):
mapping = []
for dim in e.data.subset:
syms = set()
for d in dim:
syms |= symbolic.symlist(d).keys()
if len(syms) > 1:
raise NotImplementedError("One incoming or outgoing stencil subset is indexed "
"by multiple map parameters. "
"This is not supported yet.")
try:
mapping.append(syms.pop())
except KeyError:
# just append None if there is no map symbol in it.
# we don't care for now.
mapping.append(None)
if e.data in variable_mapping:
# assert that this is the same everywhere.
# else we might run into problems
assert variable_mapping[e.data.data] == mapping
else:
variable_mapping[e.data.data] = mapping
# now do mapping data name -> outer range
# and from that infer mapping variable -> outer range
local_ranges = {dn: None for dn in coverage[map_entry][1].keys()}
for data_name, cov in coverage[map_entry][1].items():
local_ranges[data_name] = subsets.union(local_ranges[data_name], cov)
# now look at proceeding maps
# and union those subsets -> could be larger with stencil indent
for child_map in children_dict[map_entry]:
if data_name in coverage[child_map][0]:
local_ranges[data_name] = subsets.union(local_ranges[data_name],
coverage[child_map][0][data_name])
# final assignent: combine local_ranges and variable_mapping
# together into inferred_ranges
inferred_ranges[map_entry] = {p: None for p in map.params}
for data_name, ranges in local_ranges.items():
for param, r in zip(variable_mapping[data_name], ranges):
# create new range from this subset and assign
rng = subsets.Range((r, ))
if param:
inferred_ranges[map_entry][param] = subsets.union(inferred_ranges[map_entry][param], rng)
# get parameters -- should all be the same
params = next(iter(map_entries)).map.params.copy()
# define reference range as inferred range of one of the sink maps
self.reference_range = inferred_ranges[next(iter(sink_maps))]
if self.debug:
print("StencilTiling::Reference Range", self.reference_range)
# next up, search for the ranges that don't change
invariant_dims = []
for idx, p in enumerate(params):
different = False
if self.reference_range[p] is None:
invariant_dims.append(idx)
warnings.warn(f"StencilTiling::No Stencil pattern detected for parameter {p}")
continue
for m in map_entries:
if inferred_ranges[m][p] != self.reference_range[p]:
different = True
break
if not different:
invariant_dims.append(idx)
warnings.warn(f"StencilTiling::No Stencil pattern detected for parameter {p}")
# during stripmining, we will create new outer map entries
# for easy access
self._outer_entries = set()
# with inferred_ranges constructed, we can begin to strip mine
for map_entry in map_entries:
# Retrieve map entry and exit nodes.
map = map_entry.map
stripmine_subgraph = {StripMining.map_entry: graph.node_id(map_entry)}
sdfg_id = sdfg.sdfg_id
last_map_entry = None
original_schedule = map_entry.schedule
self.tile_sizes = []
self.tile_offset_lower = []
self.tile_offset_upper = []
# strip mining each dimension where necessary
removed_maps = 0
for dim_idx, param in enumerate(map_entry.map.params):
# get current_node tile size
if dim_idx >= len(self.strides):
tile_stride = symbolic.pystr_to_symbolic(self.strides[-1])
else:
tile_stride = symbolic.pystr_to_symbolic(self.strides[dim_idx])
trivial = False
if dim_idx in invariant_dims:
self.tile_sizes.append(tile_stride)
self.tile_offset_lower.append(0)
self.tile_offset_upper.append(0)
else:
target_range_current = inferred_ranges[map_entry][param]
reference_range_current = self.reference_range[param]
min_diff = symbolic.SymExpr(reference_range_current.min_element()[0] \
- target_range_current.min_element()[0])
max_diff = symbolic.SymExpr(target_range_current.max_element()[0] \
- reference_range_current.max_element()[0])
try:
min_diff = symbolic.evaluate(min_diff, {})
max_diff = symbolic.evaluate(max_diff, {})
except TypeError:
raise RuntimeError("Symbolic evaluation of map "
"ranges failed. Please check "
"your parameters and match.")
self.tile_sizes.append(tile_stride + max_diff + min_diff)
self.tile_offset_lower.append(symbolic.pystr_to_symbolic(str(min_diff)))
self.tile_offset_upper.append(symbolic.pystr_to_symbolic(str(max_diff)))
# get calculated parameters
tile_size = self.tile_sizes[-1]
dim_idx -= removed_maps
# If map or tile sizes are trivial, skip strip-mining map dimension
# special cases:
# if tile size is trivial AND we have an invariant dimension, skip
if tile_size == map.range.size()[dim_idx] and (dim_idx + removed_maps) in invariant_dims:
continue
# trivial map: we just continue
if map.range.size()[dim_idx] in [0, 1]:
continue
if tile_size == 1 and tile_stride == 1 and (dim_idx + removed_maps) in invariant_dims:
trivial = True
removed_maps += 1
# indent all map ranges accordingly and then perform
# strip mining on these. Offset inner maps accordingly afterwards
range_tuple = (map.range[dim_idx][0] + self.tile_offset_lower[-1],
map.range[dim_idx][1] - self.tile_offset_upper[-1], map.range[dim_idx][2])
map.range[dim_idx] = range_tuple
stripmine = StripMining()
stripmine.setup_match(sdfg, sdfg_id, self.state_id, stripmine_subgraph, 0)
stripmine.tiling_type = dtypes.TilingType.CeilRange
stripmine.dim_idx = dim_idx
stripmine.new_dim_prefix = self.prefix if not trivial else ''
# use tile_stride for both -- we will extend
# the inner tiles later
stripmine.tile_size = str(tile_stride)
stripmine.tile_stride = str(tile_stride)
outer_map = stripmine.apply(graph, sdfg)
outer_map.schedule = original_schedule
# if tile stride is 1, we can make a nice simplification by just
# taking the overapproximated inner range as inner range
# this eliminates the min/max in the range which
# enables loop unrolling
if not trivial:
if tile_stride == 1:
map_entry.map.range[dim_idx] = tuple(
symbolic.SymExpr(el._approx_expr) if isinstance(el, symbolic.SymExpr) else el
for el in map_entry.map.range[dim_idx])
# in map_entry: enlarge tiles by upper and lower offset
# doing it this way and not via stripmine strides ensures
# that the max gets changed as well
old_range = map_entry.map.range[dim_idx]
map_entry.map.range[dim_idx] = ((old_range[0] - self.tile_offset_lower[-1]),
(old_range[1] + self.tile_offset_upper[-1]), old_range[2])
# We have to propagate here for correct outer volume and subset sizes
_propagate_node(graph, map_entry)
_propagate_node(graph, graph.exit_node(map_entry))
# usual tiling pipeline
if last_map_entry:
new_map_entry = graph.in_edges(map_entry)[0].src
mapcollapse_subgraph = {
MapCollapse.outer_map_entry: graph.node_id(last_map_entry),
MapCollapse.inner_map_entry: graph.node_id(new_map_entry)
}
mapcollapse = MapCollapse()
mapcollapse.setup_match(sdfg, sdfg_id, self.state_id, mapcollapse_subgraph, 0)
mapcollapse.apply(graph, sdfg)
last_map_entry = graph.in_edges(map_entry)[0].src
# add last instance of map entries to _outer_entries
if last_map_entry:
self._outer_entries.add(last_map_entry)
# apply to the new map the schedule of the original one
map_entry.map.schedule = self.schedule
# Map Unroll Feature: only unroll if conditions are met:
# Only unroll if at least one of the inner map ranges is strictly larger than 1
# Only unroll if strides all are one
if self.unroll_loops and all(s == 1 for s in self.strides) and any(s not in [0, 1]
for s in map_entry.range.size()):
l = len(map_entry.params)
if l > 1:
subgraph = {MapExpansion.map_entry: graph.node_id(map_entry)}
trafo_expansion = MapExpansion()
trafo_expansion.setup_match(sdfg, sdfg.sdfg_id, sdfg.nodes().index(graph), subgraph, 0)
trafo_expansion.apply(graph, sdfg)
maps = [map_entry]
for _ in range(l - 1):
map_entry = graph.out_edges(map_entry)[0].dst
maps.append(map_entry)
for map in reversed(maps):
# MapToForLoop
subgraph = {MapToForLoop.map_entry: graph.node_id(map)}
trafo_for_loop = MapToForLoop()
trafo_for_loop.setup_match(sdfg, sdfg.sdfg_id, sdfg.nodes().index(graph), subgraph, 0)
trafo_for_loop.apply(graph, sdfg)
nsdfg = trafo_for_loop.nsdfg
# LoopUnroll
guard = trafo_for_loop.guard
end = trafo_for_loop.after_state
begin = next(e.dst for e in nsdfg.out_edges(guard) if e.dst != end)
subgraph = {
DetectLoop.loop_guard: nsdfg.node_id(guard),
DetectLoop.loop_begin: nsdfg.node_id(begin),
DetectLoop.exit_state: nsdfg.node_id(end)
}
transformation = LoopUnroll()
transformation.setup_match(nsdfg, 0, -1, subgraph, 0)
transformation.apply(nsdfg, nsdfg)
elif self.unroll_loops:
warnings.warn("StencilTiling::Did not unroll loops. Either all ranges are equal to "
"one or range difference is symbolic.")
self._outer_entries = list(self._outer_entries)
def apply(self, graph: SDFGState, sdfg: SDFG):
tile_strides = self.tile_sizes
if self.strides is not None and len(self.strides) == len(tile_strides):
tile_strides = self.strides
# Retrieve map entry and exit nodes.
map_entry = self.map_entry
from dace.transformation.dataflow.map_collapse import MapCollapse
from dace.transformation.dataflow.strip_mining import StripMining
stripmine_subgraph = {
StripMining.map_entry: self.subgraph[MapTiling.map_entry]
}
sdfg_id = sdfg.sdfg_id
last_map_entry = None
removed_maps = 0
original_schedule = map_entry.schedule
for dim_idx in range(len(map_entry.map.params)):
if dim_idx >= len(self.tile_sizes):
tile_size = symbolic.pystr_to_symbolic(self.tile_sizes[-1])
tile_stride = symbolic.pystr_to_symbolic(tile_strides[-1])
else:
tile_size = symbolic.pystr_to_symbolic(
self.tile_sizes[dim_idx])
tile_stride = symbolic.pystr_to_symbolic(tile_strides[dim_idx])
# handle offsets
if self.tile_offset and dim_idx >= len(self.tile_offset):
offset = self.tile_offset[-1]
elif self.tile_offset:
offset = self.tile_offset[dim_idx]
else:
offset = 0
dim_idx -= removed_maps
# If tile size is trivial, skip strip-mining map dimension
if tile_size == map_entry.map.range.size()[dim_idx]:
continue
stripmine = StripMining()
stripmine.setup_match(sdfg, sdfg_id, self.state_id,
stripmine_subgraph, self.expr_index)
# Special case: Tile size of 1 should be omitted from inner map
if tile_size == 1 and tile_stride == 1 and self.tile_trivial == False:
stripmine.dim_idx = dim_idx
stripmine.new_dim_prefix = ''
stripmine.tile_size = str(tile_size)
stripmine.tile_stride = str(tile_stride)
stripmine.divides_evenly = True
stripmine.tile_offset = str(offset)
stripmine.apply(graph, sdfg)
removed_maps += 1
else:
stripmine.dim_idx = dim_idx
stripmine.new_dim_prefix = self.prefix
stripmine.tile_size = str(tile_size)
stripmine.tile_stride = str(tile_stride)
stripmine.divides_evenly = self.divides_evenly
stripmine.tile_offset = str(offset)
stripmine.apply(graph, sdfg)
# apply to the new map the schedule of the original one
map_entry.schedule = original_schedule
if last_map_entry:
new_map_entry = graph.in_edges(map_entry)[0].src
mapcollapse_subgraph = {
MapCollapse.outer_map_entry: graph.node_id(last_map_entry),
MapCollapse.inner_map_entry: graph.node_id(new_map_entry)
}
mapcollapse = MapCollapse()
mapcollapse.setup_match(sdfg, sdfg_id, self.state_id,
mapcollapse_subgraph, 0)
mapcollapse.apply(graph, sdfg)
last_map_entry = graph.in_edges(map_entry)[0].src
return last_map_entry