def node_path_graph(*args):
""" Generates a path graph passing through the input nodes.
The function generates a graph using as nodes the input arguments.
Subsequently, it creates a path passing through all the nodes, in
the same order as they were given in the function input.
:param *args: Variable number of nodes or a list of nodes.
:return: A directed graph based on the input arguments.
@rtype: gr.OrderedDiGraph
"""
# 1. Create new networkx directed graph.
path = gr.OrderedDiGraph()
# 2. Place input nodes in a list.
if len(args) == 1 and isinstance(args[0], list):
# Input is a single list of nodes.
input_nodes = args[0]
else:
# Input is a variable number of nodes.
input_nodes = list(args)
# 3. Add nodes to the graph.
path.add_nodes_from(input_nodes)
# 4. Add path edges to the graph.
for i in range(len(input_nodes) - 1):
path.add_edge(input_nodes[i], input_nodes[i + 1], None)
# 5. Return the graph.
return path
class CopyToDevice(pattern_matching.Transformation):
""" Implements the copy-to-device transformation, which copies a nested
SDFG and its dependencies to a given device.
The transformation changes all data storage types of a nested SDFG to
the given `storage` property, and creates new arrays and copies around
the nested SDFG to that storage.
"""
_nested_sdfg = nodes.NestedSDFG("", graph.OrderedDiGraph(), set(), set())
storage = properties.Property(
dtype=dtypes.StorageType,
desc="Nested SDFG storage",
choices=dtypes.StorageType,
from_string=lambda x: dtypes.StorageType[x],
default=dtypes.StorageType.Default)
@staticmethod
def annotates_memlets():
return True
@staticmethod
def expressions():
return [nxutil.node_path_graph(CopyToDevice._nested_sdfg)]
@staticmethod
def can_be_applied(graph, candidate, expr_index, sdfg, strict=False):
return True
@staticmethod
def match_to_str(graph, candidate):
nested_sdfg = graph.nodes()[candidate[CopyToDevice._nested_sdfg]]
return nested_sdfg.label
def apply(self, sdfg):
state = sdfg.nodes()[self.state_id]
nested_sdfg = state.nodes()[self.subgraph[CopyToDevice._nested_sdfg]]
storage = self.storage
for _, edge in enumerate(state.in_edges(nested_sdfg)):
src, src_conn, dst, dst_conn, memlet = edge
dataname = memlet.data
memdata = sdfg.arrays[dataname]
if isinstance(memdata, data.Array):
new_data = sdfg.add_array(
'device_' + dataname + '_in',
memdata.dtype, [
symbolic.overapproximate(r)
for r in memlet.bounding_box_size()
],
transient=True,
storage=storage)
elif isinstance(memdata, data.Scalar):
new_data = sdfg.add_scalar(
'device_' + dataname + '_in',
memdata.dtype,
transient=True,
storage=storage)
else:
raise NotImplementedError
data_node = nodes.AccessNode('device_' + dataname + '_in')
to_data_mm = dcpy(memlet)
from_data_mm = dcpy(memlet)
from_data_mm.data = 'device_' + dataname + '_in'
offset = []
for ind, r in enumerate(memlet.subset):
offset.append(r[0])
if isinstance(memlet.subset[ind], tuple):
begin = memlet.subset[ind][0] - r[0]
end = memlet.subset[ind][1] - r[0]
step = memlet.subset[ind][2]
from_data_mm.subset[ind] = (begin, end, step)
else:
from_data_mm.subset[ind] -= r[0]
state.remove_edge(edge)
state.add_edge(src, src_conn, data_node, None, to_data_mm)
state.add_edge(data_node, None, dst, dst_conn, from_data_mm)
for _, edge in enumerate(state.out_edges(nested_sdfg)):
src, src_conn, dst, dst_conn, memlet = edge
dataname = memlet.data
memdata = sdfg.arrays[dataname]
if isinstance(memdata, data.Array):
new_data = data.Array(
'device_' + dataname + '_out',
memdata.dtype, [
symbolic.overapproximate(r)
for r in memlet.bounding_box_size()
],
transient=True,
storage=storage)
elif isinstance(memdata, data.Scalar):
new_data = sdfg.add_scalar(
'device_' + dataname + '_out',
memdata.dtype,
transient=True,
storage=storage)
else:
raise NotImplementedError
data_node = nodes.AccessNode('device_' + dataname + '_out')
to_data_mm = dcpy(memlet)
from_data_mm = dcpy(memlet)
to_data_mm.data = 'device_' + dataname + '_out'
offset = []
for ind, r in enumerate(memlet.subset):
offset.append(r[0])
if isinstance(memlet.subset[ind], tuple):
begin = memlet.subset[ind][0] - r[0]
end = memlet.subset[ind][1] - r[0]
step = memlet.subset[ind][2]
to_data_mm.subset[ind] = (begin, end, step)
else:
to_data_mm.subset[ind] -= r[0]
state.remove_edge(edge)
state.add_edge(src, src_conn, data_node, None, to_data_mm)
state.add_edge(data_node, None, dst, dst_conn, from_data_mm)
# Change storage for all data inside nested SDFG to device.
change_storage(nested_sdfg.sdfg, storage)