def apply(self, graph: SDFGState, sdfg: SDFG):
import dace.libraries.blas as blas
transpose_a = self.transpose_a
_at = self.at
transpose_b = self.transpose_b
_bt = self.bt
a_times_b = self.a_times_b
for src, src_conn, _, _, memlet in graph.in_edges(transpose_a):
graph.add_edge(src, src_conn, a_times_b, '_b', memlet)
graph.remove_node(transpose_a)
for src, src_conn, _, _, memlet in graph.in_edges(transpose_b):
graph.add_edge(src, src_conn, a_times_b, '_a', memlet)
graph.remove_node(transpose_b)
graph.remove_node(_at)
graph.remove_node(_bt)
for _, _, dst, dst_conn, memlet in graph.out_edges(a_times_b):
subset = dcpy(memlet.subset)
subset.squeeze()
size = subset.size()
shape = [size[1], size[0]]
break
tmp_name, tmp_arr = sdfg.add_temp_transient(shape, a_times_b.dtype)
tmp_acc = graph.add_access(tmp_name)
transpose_c = blas.Transpose('_Transpose_', a_times_b.dtype)
for edge in graph.out_edges(a_times_b):
_, _, dst, dst_conn, memlet = edge
graph.remove_edge(edge)
graph.add_edge(transpose_c, '_out', dst, dst_conn, memlet)
graph.add_edge(a_times_b, '_c', tmp_acc, None, dace.Memlet.from_array(tmp_name, tmp_arr))
graph.add_edge(tmp_acc, None, transpose_c, '_inp', dace.Memlet.from_array(tmp_name, tmp_arr))
def apply(self, graph: SDFGState, sdfg: SDFG):
tasklet = self.tasklet
map_exit = self.map_exit
outer_map_exit = self.outer_map_exit
memlet = None
edge = None
for e in graph.out_edges(map_exit):
memlet = e.data
# TODO: What if there's more than one?
if e.dst == outer_map_exit and isinstance(sdfg.arrays[memlet.data],
data.Stream):
edge = e
break
tasklet_memlet = None
for e in graph.out_edges(tasklet):
tasklet_memlet = e.data
if tasklet_memlet.data == memlet.data:
break
bbox = map_exit.map.range.bounding_box_size()
bbox_approx = [symbolic.overapproximate(dim) for dim in bbox]
dataname = memlet.data
# Create the new node: Temporary stream and an access node
newname, _ = sdfg.add_stream('trans_' + dataname,
sdfg.arrays[memlet.data].dtype,
bbox_approx[0],
storage=sdfg.arrays[memlet.data].storage,
transient=True,
find_new_name=True)
snode = graph.add_access(newname)
to_stream_mm = copy.deepcopy(memlet)
to_stream_mm.data = snode.data
tasklet_memlet.data = snode.data
if self.with_buffer:
newname_arr, _ = sdfg.add_transient('strans_' + dataname,
[bbox_approx[0]],
sdfg.arrays[memlet.data].dtype,
find_new_name=True)
anode = graph.add_access(newname_arr)
to_array_mm = copy.deepcopy(memlet)
to_array_mm.data = anode.data
graph.add_edge(snode, None, anode, None, to_array_mm)
else:
anode = snode
# Reconnect, assuming one edge to the stream
graph.remove_edge(edge)
graph.add_edge(map_exit, edge.src_conn, snode, None, to_stream_mm)
graph.add_edge(anode, None, outer_map_exit, edge.dst_conn, memlet)
return
def memlets_intersect(graph_a: SDFGState, group_a: List[nodes.AccessNode],
inputs_a: bool, graph_b: SDFGState,
group_b: List[nodes.AccessNode],
inputs_b: bool) -> bool:
"""
Performs an all-pairs check for subset intersection on two
groups of nodes. If group intersects or result is indeterminate,
returns True as a precaution.
:param graph_a: The graph in which the first set of nodes reside.
:param group_a: The first set of nodes to check.
:param inputs_a: If True, checks inputs of the first group.
:param graph_b: The graph in which the second set of nodes reside.
:param group_b: The second set of nodes to check.
:param inputs_b: If True, checks inputs of the second group.
:returns True if subsets intersect or result is indeterminate.
"""
# Set traversal functions
src_subset = lambda e: (e.data.src_subset if e.data.src_subset is
not None else e.data.dst_subset)
dst_subset = lambda e: (e.data.dst_subset if e.data.dst_subset is
not None else e.data.src_subset)
if inputs_a:
edges_a = [e for n in group_a for e in graph_a.out_edges(n)]
subset_a = src_subset
else:
edges_a = [e for n in group_a for e in graph_a.in_edges(n)]
subset_a = dst_subset
if inputs_b:
edges_b = [e for n in group_b for e in graph_b.out_edges(n)]
subset_b = src_subset
else:
edges_b = [e for n in group_b for e in graph_b.in_edges(n)]
subset_b = dst_subset
# Simple all-pairs check
for ea in edges_a:
for eb in edges_b:
result = subsets.intersects(subset_a(ea), subset_b(eb))
if result is True or result is None:
return True
return False
def remove_edge_and_dangling_path(state: SDFGState, edge: MultiConnectorEdge):
"""
Removes an edge and all of its parent edges in a memlet path, cleaning
dangling connectors and isolated nodes resulting from the removal.
:param state: The state in which the edge exists.
:param edge: The edge to remove.
"""
mtree = state.memlet_tree(edge)
inwards = (isinstance(edge.src, nd.EntryNode)
or isinstance(edge.dst, nd.EntryNode))
# Traverse tree upwards, removing edges and connectors as necessary
curedge = mtree
while curedge is not None:
e = curedge.edge
state.remove_edge(e)
if inwards:
neighbors = [
neighbor for neighbor in state.out_edges(e.src)
if e.src_conn == neighbor.src_conn
]
else:
neighbors = [
neighbor for neighbor in state.in_edges(e.dst)
if e.dst_conn == neighbor.dst_conn
]
if len(neighbors) > 0: # There are still edges connected, leave as-is
break
# Remove connector and matching outer connector
if inwards:
if e.src_conn:
e.src.remove_out_connector(e.src_conn)
e.src.remove_in_connector('IN' + e.src_conn[3:])
else:
if e.dst_conn:
e.dst.remove_in_connector(e.dst_conn)
e.src.remove_out_connector('OUT' + e.dst_conn[2:])
# Continue traversing upwards
curedge = curedge.parent
else:
# Check if an isolated node have been created at the root and remove
root_edge = mtree.root().edge
root_node: nd.Node = root_edge.src if inwards else root_edge.dst
if state.degree(root_node) == 0:
state.remove_node(root_node)
def is_array_stream_view(sdfg: SDFG, dfg: SDFGState, node: nd.AccessNode):
""" Test whether a stream is directly connected to an array. """
# Test all memlet paths from the array. If the path goes directly
# to/from a stream, construct a stream array view
all_source_paths = []
source_paths = []
all_sink_paths = []
sink_paths = []
for e in dfg.in_edges(node):
src_node = dfg.memlet_path(e)[0].src
# Append empty path to differentiate between a copy and an array-view
if isinstance(src_node, nd.CodeNode):
all_source_paths.append(None)
# Append path from source node
if isinstance(src_node, nd.AccessNode) and isinstance(
src_node.desc(sdfg), dt.Array):
source_paths.append(src_node)
for e in dfg.out_edges(node):
sink_node = dfg.memlet_path(e)[-1].dst
# Append empty path to differentiate between a copy and an array-view
if isinstance(sink_node, nd.CodeNode):
all_sink_paths.append(None)
# Append path to sink node
if isinstance(sink_node, nd.AccessNode) and isinstance(
sink_node.desc(sdfg), dt.Array):
sink_paths.append(sink_node)
all_sink_paths.extend(sink_paths)
all_source_paths.extend(source_paths)
# Special case: stream can be represented as a view of an array
if ((len(all_source_paths) > 0 and len(sink_paths) == 1)
or (len(all_sink_paths) > 0 and len(source_paths) == 1)):
# TODO: What about a source path?
arrnode = sink_paths[0]
# Only works if the stream itself is not an array of streams
if list(node.desc(sdfg).shape) == [1]:
node.desc(sdfg).sink = arrnode.data # For memlet generation
arrnode.desc(
sdfg).src = node.data # TODO: Move src/sink to node, not array
return True
return False
def apply(self, graph: SDFGState, sdfg: SDFG):
first_map_entry = graph.entry_node(self.first_map_exit)
intermediate_dnodes = set()
for _, _, node, _, _ in graph.out_edges(self.first_map_exit):
if not isinstance(node, nds.AccessNode):
continue
intermediate_dnodes.add(node)
self._update_in_connectors(graph, intermediate_dnodes)
self._replicate_first_map(sdfg, graph, first_map_entry,
intermediate_dnodes)
graph.remove_nodes_from(
graph.all_nodes_between(first_map_entry, self.first_map_exit)
| {first_map_entry, self.first_map_exit})
for node in graph.nodes():
if not isinstance(node, nds.AccessNode):
continue
if graph.in_degree(node) == 0 and graph.out_degree(node) == 0:
graph.remove_node(node)
def trace_nested_access(
node: nd.AccessNode, state: SDFGState,
sdfg: SDFG) -> List[Tuple[nd.AccessNode, SDFGState, SDFG]]:
"""
Given an AccessNode in a nested SDFG, trace the accessed memory
back to the outermost scope in which it is defined.
:param node: An access node.
:param state: State in which the access node is located.
:param sdfg: SDFG in which the access node is located.
:return: A list of scopes ((input_node, output_node), (memlet_read, memlet_write), state, sdfg) in which
the given data is accessed, from outermost scope to innermost
scope.
"""
curr_sdfg = sdfg
curr_read = None
memlet_read = None
for m in state.out_edges(node):
if not m.data.is_empty():
curr_read = node
memlet_read = m.data
break
curr_write = None
memlet_write = None
for m in state.in_edges(node):
if not m.data.is_empty():
curr_write = node
memlet_read = m.data
break
trace = [((curr_read, curr_write), (memlet_read, memlet_write), state, sdfg)
]
while curr_sdfg.parent is not None:
curr_state = curr_sdfg.parent
# Find the nested SDFG containing ourself in the parent state
nested_sdfg = curr_sdfg.parent_nsdfg_node
if curr_read is not None:
for e in curr_state.in_edges(nested_sdfg):
if e.dst_conn == curr_read.data:
# See if the input to this connector traces back to an
# access node. If not, just give up here
n = find_input_arraynode(curr_state, e)
if isinstance(n, nd.AccessNode):
curr_read = n
memlet_read = e.data
break
else:
curr_read = None
memlet_read = None
if curr_write is not None:
for e in curr_state.out_edges(nested_sdfg):
if e.src_conn == curr_write.data:
# See if the output of this connector traces back to an
# access node. If not, just give up here
n = find_output_arraynode(curr_state, e)
if isinstance(curr_write, nd.AccessNode):
curr_write = n
memlet_write = e.data
break
else:
curr_write = None
memlet_write = None
if curr_read is not None or curr_write is not None:
trace.append(((curr_read, curr_write), (memlet_read, memlet_write),
curr_state, curr_state.parent))
else:
break
curr_sdfg = curr_state.parent # Recurse
return list(reversed(trace))
def get_view_edge(state: SDFGState,
view: nd.AccessNode) -> gr.MultiConnectorEdge[mm.Memlet]:
"""
Given a view access node, returns the
incoming/outgoing edge which points to the viewed access node.
See the ruleset in the documentation of ``dace.data.View``.
:param state: The state in which the view resides.
:param view: The view access node.
:return: An edge pointing to the viewed data or None if view is invalid.
:see: ``dace.data.View``
"""
in_edges = state.in_edges(view)
out_edges = state.out_edges(view)
# Invalid case: No data to view
if len(in_edges) == 0 or len(out_edges) == 0:
return None
# If there is one edge (in/out) that leads (via memlet path) to an access
# node, and the other side (out/in) has a different number of edges.
if len(in_edges) == 1 and len(out_edges) != 1:
return in_edges[0]
if len(out_edges) == 1 and len(in_edges) != 1:
return out_edges[0]
if len(out_edges) == len(in_edges) and len(out_edges) != 1:
return None
in_edge = in_edges[0]
out_edge = out_edges[0]
# If there is one incoming and one outgoing edge, and one leads to a code
# node, the one that leads to an access node is the viewed data.
inmpath = state.memlet_path(in_edge)
outmpath = state.memlet_path(out_edge)
src_is_data, dst_is_data = False, False
if isinstance(inmpath[0].src, nd.AccessNode):
src_is_data = True
if isinstance(outmpath[-1].dst, nd.AccessNode):
dst_is_data = True
if src_is_data and not dst_is_data:
return in_edge
if not src_is_data and dst_is_data:
return out_edge
if not src_is_data and not dst_is_data:
return None
# If both sides lead to access nodes, if one memlet's data points to the
# view it cannot point to the viewed node.
if in_edge.data.data == view.data and out_edge.data.data != view.data:
return out_edge
if in_edge.data.data != view.data and out_edge.data.data == view.data:
return in_edge
if in_edge.data.data == view.data and out_edge.data.data == view.data:
return None
# If both memlets' data are the respective access nodes, the access
# node at the highest scope is the one that is viewed.
if isinstance(in_edge.src, nd.EntryNode):
return in_edge
if isinstance(out_edge.dst, nd.ExitNode):
return out_edge
# If both access nodes reside in the same scope, the input data is viewed.
return in_edge
def merge_maps(
graph: SDFGState,
outer_map_entry: nd.MapEntry,
outer_map_exit: nd.MapExit,
inner_map_entry: nd.MapEntry,
inner_map_exit: nd.MapExit,
param_merge: Callable[[ParamsType, ParamsType],
ParamsType] = lambda p1, p2: p1 + p2,
range_merge: Callable[[RangesType, RangesType],
RangesType] = lambda r1, r2: type(r1)
(r1.ranges + r2.ranges)
) -> (nd.MapEntry, nd.MapExit):
""" Merges two maps (their entries and exits). It is assumed that the
operation is valid. """
outer_map = outer_map_entry.map
inner_map = inner_map_entry.map
# Create merged map by inheriting attributes from outer map and using
# the merge functions for parameters and ranges.
merged_map = copy.deepcopy(outer_map)
merged_map.label = outer_map.label
merged_map.params = param_merge(outer_map.params, inner_map.params)
merged_map.range = range_merge(outer_map.range, inner_map.range)
merged_entry = nd.MapEntry(merged_map)
merged_entry.in_connectors = outer_map_entry.in_connectors
merged_entry.out_connectors = outer_map_entry.out_connectors
merged_exit = nd.MapExit(merged_map)
merged_exit.in_connectors = outer_map_exit.in_connectors
merged_exit.out_connectors = outer_map_exit.out_connectors
graph.add_nodes_from([merged_entry, merged_exit])
# Handle the case of dynamic map inputs in the inner map
inner_dynamic_map_inputs = dynamic_map_inputs(graph, inner_map_entry)
for edge in inner_dynamic_map_inputs:
remove_conn = (len(
list(graph.out_edges_by_connector(edge.src, edge.src_conn))) == 1)
conn_to_remove = edge.src_conn[4:]
if remove_conn:
merged_entry.remove_in_connector('IN_' + conn_to_remove)
merged_entry.remove_out_connector('OUT_' + conn_to_remove)
merged_entry.add_in_connector(
edge.dst_conn, inner_map_entry.in_connectors[edge.dst_conn])
outer_edge = next(
graph.in_edges_by_connector(outer_map_entry,
'IN_' + conn_to_remove))
graph.add_edge(outer_edge.src, outer_edge.src_conn, merged_entry,
edge.dst_conn, outer_edge.data)
if remove_conn:
graph.remove_edge(outer_edge)
# Redirect inner in edges.
for edge in graph.out_edges(inner_map_entry):
if edge.src_conn is None: # Empty memlets
graph.add_edge(merged_entry, edge.src_conn, edge.dst, edge.dst_conn,
edge.data)
continue
# Get memlet path and edge
path = graph.memlet_path(edge)
ind = path.index(edge)
# Add an edge directly from the previous source connector to the
# destination
graph.add_edge(merged_entry, path[ind - 1].src_conn, edge.dst,
edge.dst_conn, edge.data)
# Redirect inner out edges.
for edge in graph.in_edges(inner_map_exit):
if edge.dst_conn is None: # Empty memlets
graph.add_edge(edge.src, edge.src_conn, merged_exit, edge.dst_conn,
edge.data)
continue
# Get memlet path and edge
path = graph.memlet_path(edge)
ind = path.index(edge)
# Add an edge directly from the source to the next destination
# connector
graph.add_edge(edge.src, edge.src_conn, merged_exit,
path[ind + 1].dst_conn, edge.data)
# Redirect outer edges.
change_edge_dest(graph, outer_map_entry, merged_entry)
change_edge_src(graph, outer_map_exit, merged_exit)
# Clean-up
graph.remove_nodes_from(
[outer_map_entry, outer_map_exit, inner_map_entry, inner_map_exit])
return merged_entry, merged_exit
def apply(self, graph: SDFGState, sdfg: SDFG):
"""
This method applies the mapfusion transformation.
Other than the removal of the second map entry node (SME), and the first
map exit (FME) node, it has the following side effects:
1. Any transient adjacent to both FME and SME with degree = 2 will be removed.
The tasklets that use/produce it shall be connected directly with a
scalar/new transient (if the dataflow is more than a single scalar)
2. If this transient is adjacent to FME and SME and has other
uses, it will be adjacent to the new map exit post fusion.
Tasklet-> Tasklet edges will ALSO be added as mentioned above.
3. If an access node is adjacent to FME but not SME, it will be
adjacent to new map exit post fusion.
4. If an access node is adjacent to SME but not FME, it will be
adjacent to the new map entry node post fusion.
"""
first_exit = self.first_map_exit
first_entry = graph.entry_node(first_exit)
second_entry = self.second_map_entry
second_exit = graph.exit_node(second_entry)
intermediate_nodes = set()
for _, _, dst, _, _ in graph.out_edges(first_exit):
intermediate_nodes.add(dst)
assert isinstance(dst, nodes.AccessNode)
# Check if an access node refers to non transient memory, or transient
# is used at another location (cannot erase)
do_not_erase = set()
for node in intermediate_nodes:
if sdfg.arrays[node.data].transient is False:
do_not_erase.add(node)
else:
for edge in graph.in_edges(node):
if edge.src != first_exit:
do_not_erase.add(node)
break
else:
for edge in graph.out_edges(node):
if edge.dst != second_entry:
do_not_erase.add(node)
break
# Find permutation between first and second scopes
perm = self.find_permutation(first_entry.map, second_entry.map)
params_dict = {}
for index, param in enumerate(first_entry.map.params):
params_dict[param] = second_entry.map.params[perm[index]]
# Replaces (in memlets and tasklet) the second scope map
# indices with the permuted first map indices.
# This works in two passes to avoid problems when e.g., exchanging two
# parameters (instead of replacing (j,i) and (i,j) to (j,j) and then
# i,i).
second_scope = graph.scope_subgraph(second_entry)
for firstp, secondp in params_dict.items():
if firstp != secondp:
replace(second_scope, secondp, '__' + secondp + '_fused')
for firstp, secondp in params_dict.items():
if firstp != secondp:
replace(second_scope, '__' + secondp + '_fused', firstp)
# Isolate First exit node
############################
edges_to_remove = set()
nodes_to_remove = set()
for edge in graph.in_edges(first_exit):
tree = graph.memlet_tree(edge)
access_node = tree.root().edge.dst
if access_node not in do_not_erase:
out_edges = [
e for e in graph.out_edges(access_node)
if e.dst == second_entry
]
# In this transformation, there can only be one edge to the
# second map
assert len(out_edges) == 1
# Get source connector to the second map
connector = out_edges[0].dst_conn[3:]
new_dsts = []
# Look at the second map entry out-edges to get the new
# destinations
for e in graph.out_edges(second_entry):
if e.src_conn[4:] == connector:
new_dsts.append(e)
if not new_dsts: # Access node is not used in the second map
nodes_to_remove.add(access_node)
continue
# Add a transient scalar/array
self.fuse_nodes(sdfg, graph, edge, new_dsts[0].dst,
new_dsts[0].dst_conn, new_dsts[1:])
edges_to_remove.add(edge)
# Remove transient node between the two maps
nodes_to_remove.add(access_node)
else: # The case where intermediate array node cannot be removed
# Node will become an output of the second map exit
out_e = tree.parent.edge
conn = second_exit.next_connector()
graph.add_edge(
second_exit,
'OUT_' + conn,
out_e.dst,
out_e.dst_conn,
dcpy(out_e.data),
)
second_exit.add_out_connector('OUT_' + conn)
graph.add_edge(edge.src, edge.src_conn, second_exit,
'IN_' + conn, dcpy(edge.data))
second_exit.add_in_connector('IN_' + conn)
edges_to_remove.add(out_e)
edges_to_remove.add(edge)
# If the second map needs this node, link the connector
# that generated this to the place where it is needed, with a
# temp transient/scalar for memlet to be generated
for out_e in graph.out_edges(second_entry):
second_memlet_path = graph.memlet_path(out_e)
source_node = second_memlet_path[0].src
if source_node == access_node:
self.fuse_nodes(sdfg, graph, edge, out_e.dst,
out_e.dst_conn)
###
# First scope exit is isolated and can now be safely removed
for e in edges_to_remove:
graph.remove_edge(e)
graph.remove_nodes_from(nodes_to_remove)
graph.remove_node(first_exit)
# Isolate second_entry node
###########################
for edge in graph.in_edges(second_entry):
tree = graph.memlet_tree(edge)
access_node = tree.root().edge.src
if access_node in intermediate_nodes:
# Already handled above, can be safely removed
graph.remove_edge(edge)
continue
# This is an external input to the second map which will now go
# through the first map.
conn = first_entry.next_connector()
graph.add_edge(edge.src, edge.src_conn, first_entry, 'IN_' + conn,
dcpy(edge.data))
first_entry.add_in_connector('IN_' + conn)
graph.remove_edge(edge)
for out_enode in tree.children:
out_e = out_enode.edge
graph.add_edge(
first_entry,
'OUT_' + conn,
out_e.dst,
out_e.dst_conn,
dcpy(out_e.data),
)
graph.remove_edge(out_e)
first_entry.add_out_connector('OUT_' + conn)
###
# Second node is isolated and can now be safely removed
graph.remove_node(second_entry)
# Fix scope exit to point to the right map
second_exit.map = first_entry.map
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(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:
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(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):
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(sdfg, sdfg_id, self.state_id,
outlocalstorage_subgraph,
self.expr_index)
outlocalstorage.array = name
outlocalstorage.apply(graph, sdfg)
def infer_node_connectors(sdfg: SDFG, state: SDFGState, node: nodes.Node,
inferred: TypeInferenceDict):
""" Infers the connectors of a node and updates `inferred` accordingly. """
# Try to infer input connector type from node type or previous edges
for e in state.in_edges(node):
cname = e.dst_conn
if cname is None:
continue
scalar = (e.data.subset and e.data.subset.num_elements() == 1)
if e.data.data is not None:
allocated_as_scalar = (sdfg.arrays[e.data.data].storage
is not dtypes.StorageType.GPU_Global)
else:
allocated_as_scalar = True
if inferred[(node, cname, True)].type is None:
# If nested SDFG, try to use internal array type
if isinstance(node, nodes.NestedSDFG):
scalar = (isinstance(node.sdfg.arrays[cname], data.Scalar)
and allocated_as_scalar)
dtype = node.sdfg.arrays[cname].dtype
ctype = (dtype if scalar else dtypes.pointer(dtype))
elif e.data.data is not None: # Obtain type from memlet
scalar |= isinstance(sdfg.arrays[e.data.data], data.Scalar)
if isinstance(node, nodes.LibraryNode):
scalar &= allocated_as_scalar
dtype = sdfg.arrays[e.data.data].dtype
ctype = (dtype if scalar else dtypes.pointer(dtype))
else: # Code->Code
src_edge = state.memlet_path(e)[0]
sconn = src_edge.src.out_connectors[src_edge.src_conn]
if sconn.type is None:
raise TypeError('Ambiguous or uninferable type in'
' connector "%s" of node "%s"' %
(sconn, src_edge.src))
ctype = sconn
inferred[(node, cname, True)] = ctype
# Try to infer outputs from output edges
for e in state.out_edges(node):
cname = e.src_conn
if cname is None:
continue
scalar = (e.data.subset and e.data.subset.num_elements() == 1
and (not e.data.dynamic or
(e.data.dynamic and e.data.wcr is not None)))
if e.data.data is not None:
allocated_as_scalar = (sdfg.arrays[e.data.data].storage
is not dtypes.StorageType.GPU_Global)
else:
allocated_as_scalar = True
if inferred[(node, cname, False)].type is None:
# If nested SDFG, try to use internal array type
if isinstance(node, nodes.NestedSDFG):
scalar = (isinstance(node.sdfg.arrays[cname], data.Scalar)
and allocated_as_scalar)
dtype = node.sdfg.arrays[cname].dtype
ctype = (dtype if scalar else dtypes.pointer(dtype))
elif e.data.data is not None: # Obtain type from memlet
scalar |= isinstance(sdfg.arrays[e.data.data], data.Scalar)
if isinstance(node, nodes.LibraryNode):
scalar &= allocated_as_scalar
dtype = sdfg.arrays[e.data.data].dtype
ctype = (dtype if scalar else dtypes.pointer(dtype))
else:
continue
inferred[(node, cname, False)] = ctype
# Let the node infer other output types on its own
if isinstance(node, nodes.Tasklet):
infer_tasklet_connectors(sdfg, state, node, inferred)
elif isinstance(node, nodes.NestedSDFG):
infer_connector_types(node.sdfg, inferred=inferred)
# If there are any remaining uninferable connectors, fail
for e in state.out_edges(node):
cname = e.src_conn
if cname and inferred[(node, cname, False)].type is None:
raise TypeError('Ambiguous or uninferable type in'
' connector "%s" of node "%s"' % (cname, node))
