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
def expressions(cls):
graph = gr.OrderedDiGraph()
graph.add_node(cls.transpose_a)
graph.add_node(cls.at)
graph.add_node(cls.transpose_b)
graph.add_node(cls.bt)
graph.add_node(cls.a_times_b)
graph.add_edge(cls.transpose_a, cls.at, None)
graph.add_edge(cls.at, cls.a_times_b, None)
graph.add_edge(cls.transpose_b, cls.bt, None)
graph.add_edge(cls.bt, cls.a_times_b, None)
return [graph]
def _make_dependency_graph(self) -> gr.OrderedDiGraph:
"""
Makes an ordered dependency graph out of the passes in the pipeline to traverse when applying.
"""
result = gr.OrderedDiGraph()
ptype_to_pass = {type(p): p for p in self.passes}
for p in self.passes:
if p not in result._nodes:
result.add_node(p)
for dep in p.depends_on():
# If a type, find it in self.passes
if isinstance(dep, type):
dep = ptype_to_pass[dep]
result.add_edge(dep, p)
return result
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(), {}, {})
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 [sdutil.node_path_graph(CopyToDevice._nested_sdfg)]
@staticmethod
def can_be_applied(graph, candidate, expr_index, sdfg, strict=False):
nested_sdfg = graph.nodes()[candidate[CopyToDevice._nested_sdfg]]
for edge in graph.all_edges(nested_sdfg):
# Stream inputs/outputs not allowed
path = graph.memlet_path(edge)
if ((isinstance(path[0].src, nodes.AccessNode)
and isinstance(sdfg.arrays[path[0].src.data], data.Stream)) or
(isinstance(path[-1].dst, nodes.AccessNode)
and isinstance(sdfg.arrays[path[-1].dst.data], data.Stream))):
return False
# WCR outputs with arrays are not allowed
if (edge.data.wcr is not None
and edge.data.subset.num_elements() != 1):
return 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
created_arrays = set()
for _, edge in enumerate(state.in_edges(nested_sdfg)):
src, src_conn, dst, dst_conn, memlet = edge
dataname = memlet.data
if dataname is None:
continue
memdata = sdfg.arrays[dataname]
name = 'device_' + dataname + '_in'
if name not in created_arrays:
if isinstance(memdata, data.Array):
name, _ = sdfg.add_array(
'device_' + dataname + '_in',
shape=[
symbolic.overapproximate(r)
for r in memlet.bounding_box_size()
],
dtype=memdata.dtype,
transient=True,
storage=storage,
find_new_name=True)
elif isinstance(memdata, data.Scalar):
name, _ = sdfg.add_scalar('device_' + dataname + '_in',
dtype=memdata.dtype,
transient=True,
storage=storage,
find_new_name=True)
else:
raise NotImplementedError
created_arrays.add(name)
data_node = nodes.AccessNode(name)
to_data_mm = dcpy(memlet)
from_data_mm = dcpy(memlet)
from_data_mm.data = name
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
if dataname is None:
continue
memdata = sdfg.arrays[dataname]
name = 'device_' + dataname + '_out'
if name not in created_arrays:
if isinstance(memdata, data.Array):
name, _ = sdfg.add_array(
name,
shape=[
symbolic.overapproximate(r)
for r in memlet.bounding_box_size()
],
dtype=memdata.dtype,
transient=True,
storage=storage,
find_new_name=True)
elif isinstance(memdata, data.Scalar):
name, _ = sdfg.add_scalar(name,
dtype=memdata.dtype,
transient=True,
storage=storage)
else:
raise NotImplementedError
created_arrays.add(name)
data_node = nodes.AccessNode(name)
to_data_mm = dcpy(memlet)
from_data_mm = dcpy(memlet)
to_data_mm.data = name
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)
