def nest_state_subgraph(sdfg: SDFG,
state: SDFGState,
subgraph: SubgraphView,
name: Optional[str] = None,
full_data: bool = False) -> nodes.NestedSDFG:
""" Turns a state subgraph into a nested SDFG. Operates in-place.
:param sdfg: The SDFG containing the state subgraph.
:param state: The state containing the subgraph.
:param subgraph: Subgraph to nest.
:param name: An optional name for the nested SDFG.
:param full_data: If True, nests entire input/output data.
:return: The nested SDFG node.
:raise KeyError: Some or all nodes in the subgraph are not located in
this state, or the state does not belong to the given
SDFG.
:raise ValueError: The subgraph is contained in more than one scope.
"""
if state.parent != sdfg:
raise KeyError('State does not belong to given SDFG')
if subgraph.graph != state:
raise KeyError('Subgraph does not belong to given state')
# Find the top-level scope
scope_tree = state.scope_tree()
scope_dict = state.scope_dict()
scope_dict_children = state.scope_dict(True)
top_scopenode = -1 # Initialized to -1 since "None" already means top-level
for node in subgraph.nodes():
if node not in scope_dict:
raise KeyError('Node not found in state')
# If scope entry/exit, ensure entire scope is in subgraph
if isinstance(node, nodes.EntryNode):
scope_nodes = scope_dict_children[node]
if any(n not in subgraph.nodes() for n in scope_nodes):
raise ValueError('Subgraph contains partial scopes (entry)')
elif isinstance(node, nodes.ExitNode):
entry = state.entry_node(node)
scope_nodes = scope_dict_children[entry] + [entry]
if any(n not in subgraph.nodes() for n in scope_nodes):
raise ValueError('Subgraph contains partial scopes (exit)')
scope_node = scope_dict[node]
if scope_node not in subgraph.nodes():
if top_scopenode != -1 and top_scopenode != scope_node:
raise ValueError(
'Subgraph is contained in more than one scope')
top_scopenode = scope_node
scope = scope_tree[top_scopenode]
###
# Collect inputs and outputs of the nested SDFG
inputs: List[MultiConnectorEdge] = []
outputs: List[MultiConnectorEdge] = []
for node in subgraph.source_nodes():
inputs.extend(state.in_edges(node))
for node in subgraph.sink_nodes():
outputs.extend(state.out_edges(node))
# Collect transients not used outside of subgraph (will be removed of
# top-level graph)
data_in_subgraph = set(n.data for n in subgraph.nodes()
if isinstance(n, nodes.AccessNode))
# Find other occurrences in SDFG
other_nodes = set(
n.data for s in sdfg.nodes() for n in s.nodes()
if isinstance(n, nodes.AccessNode) and n not in subgraph.nodes())
subgraph_transients = set()
for data in data_in_subgraph:
datadesc = sdfg.arrays[data]
if datadesc.transient and data not in other_nodes:
subgraph_transients.add(data)
# All transients of edges between code nodes are also added to nested graph
for edge in subgraph.edges():
if (isinstance(edge.src, nodes.CodeNode)
and isinstance(edge.dst, nodes.CodeNode)):
subgraph_transients.add(edge.data.data)
# Collect data used in access nodes within subgraph (will be referenced in
# full upon nesting)
input_arrays = set()
output_arrays = set()
for node in subgraph.nodes():
if (isinstance(node, nodes.AccessNode)
and node.data not in subgraph_transients):
if state.out_degree(node) > 0:
input_arrays.add(node.data)
if state.in_degree(node) > 0:
output_arrays.add(node.data)
# Create the nested SDFG
nsdfg = SDFG(name or 'nested_' + state.label)
# Transients are added to the nested graph as-is
for name in subgraph_transients:
nsdfg.add_datadesc(name, sdfg.arrays[name])
# Input/output data that are not source/sink nodes are added to the graph
# as non-transients
for name in (input_arrays | output_arrays):
datadesc = copy.deepcopy(sdfg.arrays[name])
datadesc.transient = False
nsdfg.add_datadesc(name, datadesc)
# Connected source/sink nodes outside subgraph become global data
# descriptors in nested SDFG
input_names = []
output_names = []
for edge in inputs:
if edge.data.data is None: # Skip edges with an empty memlet
continue
name = '__in_' + edge.data.data
datadesc = copy.deepcopy(sdfg.arrays[edge.data.data])
datadesc.transient = False
if not full_data:
datadesc.shape = edge.data.subset.size()
input_names.append(
nsdfg.add_datadesc(name, datadesc, find_new_name=True))
for edge in outputs:
if edge.data.data is None: # Skip edges with an empty memlet
continue
name = '__out_' + edge.data.data
datadesc = copy.deepcopy(sdfg.arrays[edge.data.data])
datadesc.transient = False
if not full_data:
datadesc.shape = edge.data.subset.size()
output_names.append(
nsdfg.add_datadesc(name, datadesc, find_new_name=True))
###################
# Add scope symbols to the nested SDFG
for v in scope.defined_vars:
if v in sdfg.symbols:
sym = sdfg.symbols[v]
nsdfg.add_symbol(v, sym.dtype)
# Create nested state
nstate = nsdfg.add_state()
# Add subgraph nodes and edges to nested state
nstate.add_nodes_from(subgraph.nodes())
for e in subgraph.edges():
nstate.add_edge(e.src, e.src_conn, e.dst, e.dst_conn, e.data)
# Modify nested SDFG parents in subgraph
for node in subgraph.nodes():
if isinstance(node, nodes.NestedSDFG):
node.sdfg.parent = nstate
node.sdfg.parent_sdfg = nsdfg
# Add access nodes and edges as necessary
edges_to_offset = []
for name, edge in zip(input_names, inputs):
node = nstate.add_read(name)
new_edge = copy.deepcopy(edge.data)
new_edge.data = name
edges_to_offset.append((edge,
nstate.add_edge(node, None, edge.dst,
edge.dst_conn, new_edge)))
for name, edge in zip(output_names, outputs):
node = nstate.add_write(name)
new_edge = copy.deepcopy(edge.data)
new_edge.data = name
edges_to_offset.append((edge,
nstate.add_edge(edge.src, edge.src_conn, node,
None, new_edge)))
# Offset memlet paths inside nested SDFG according to subsets
for original_edge, new_edge in edges_to_offset:
for edge in nstate.memlet_tree(new_edge):
edge.data.data = new_edge.data.data
if not full_data:
edge.data.subset.offset(original_edge.data.subset, True)
# Add nested SDFG node to the input state
nested_sdfg = state.add_nested_sdfg(nsdfg, None,
set(input_names) | input_arrays,
set(output_names) | output_arrays)
# Reconnect memlets to nested SDFG
for name, edge in zip(input_names, inputs):
if full_data:
data = Memlet.from_array(edge.data.data,
sdfg.arrays[edge.data.data])
else:
data = edge.data
state.add_edge(edge.src, edge.src_conn, nested_sdfg, name, data)
for name, edge in zip(output_names, outputs):
if full_data:
data = Memlet.from_array(edge.data.data,
sdfg.arrays[edge.data.data])
else:
data = edge.data
state.add_edge(nested_sdfg, name, edge.dst, edge.dst_conn, data)
# Connect access nodes to internal input/output data as necessary
entry = scope.entry
exit = scope.exit
for name in input_arrays:
node = state.add_read(name)
if entry is not None:
state.add_nedge(entry, node, EmptyMemlet())
state.add_edge(node, None, nested_sdfg, name,
Memlet.from_array(name, sdfg.arrays[name]))
for name in output_arrays:
node = state.add_write(name)
if exit is not None:
state.add_nedge(node, exit, EmptyMemlet())
state.add_edge(nested_sdfg, name, node, None,
Memlet.from_array(name, sdfg.arrays[name]))
# Remove subgraph nodes from graph
state.remove_nodes_from(subgraph.nodes())
# Remove subgraph transients from top-level graph
for transient in subgraph_transients:
del sdfg.arrays[transient]
return nested_sdfg
def nest_state_subgraph(sdfg: SDFG,
state: SDFGState,
subgraph: SubgraphView,
name: Optional[str] = None,
full_data: bool = False) -> nodes.NestedSDFG:
""" Turns a state subgraph into a nested SDFG. Operates in-place.
:param sdfg: The SDFG containing the state subgraph.
:param state: The state containing the subgraph.
:param subgraph: Subgraph to nest.
:param name: An optional name for the nested SDFG.
:param full_data: If True, nests entire input/output data.
:return: The nested SDFG node.
:raise KeyError: Some or all nodes in the subgraph are not located in
this state, or the state does not belong to the given
SDFG.
:raise ValueError: The subgraph is contained in more than one scope.
"""
if state.parent != sdfg:
raise KeyError('State does not belong to given SDFG')
if subgraph is not state and subgraph.graph is not state:
raise KeyError('Subgraph does not belong to given state')
# Find the top-level scope
scope_tree = state.scope_tree()
scope_dict = state.scope_dict()
scope_dict_children = state.scope_children()
top_scopenode = -1 # Initialized to -1 since "None" already means top-level
for node in subgraph.nodes():
if node not in scope_dict:
raise KeyError('Node not found in state')
# If scope entry/exit, ensure entire scope is in subgraph
if isinstance(node, nodes.EntryNode):
scope_nodes = scope_dict_children[node]
if any(n not in subgraph.nodes() for n in scope_nodes):
raise ValueError('Subgraph contains partial scopes (entry)')
elif isinstance(node, nodes.ExitNode):
entry = state.entry_node(node)
scope_nodes = scope_dict_children[entry] + [entry]
if any(n not in subgraph.nodes() for n in scope_nodes):
raise ValueError('Subgraph contains partial scopes (exit)')
scope_node = scope_dict[node]
if scope_node not in subgraph.nodes():
if top_scopenode != -1 and top_scopenode != scope_node:
raise ValueError('Subgraph is contained in more than one scope')
top_scopenode = scope_node
scope = scope_tree[top_scopenode]
###
# Consolidate edges in top scope
utils.consolidate_edges(sdfg, scope)
snodes = subgraph.nodes()
# Collect inputs and outputs of the nested SDFG
inputs: List[MultiConnectorEdge] = []
outputs: List[MultiConnectorEdge] = []
for node in snodes:
for edge in state.in_edges(node):
if edge.src not in snodes:
inputs.append(edge)
for edge in state.out_edges(node):
if edge.dst not in snodes:
outputs.append(edge)
# Collect transients not used outside of subgraph (will be removed of
# top-level graph)
data_in_subgraph = set(n.data for n in subgraph.nodes() if isinstance(n, nodes.AccessNode))
# Find other occurrences in SDFG
other_nodes = set(n.data for s in sdfg.nodes() for n in s.nodes()
if isinstance(n, nodes.AccessNode) and n not in subgraph.nodes())
subgraph_transients = set()
for data in data_in_subgraph:
datadesc = sdfg.arrays[data]
if datadesc.transient and data not in other_nodes:
subgraph_transients.add(data)
# All transients of edges between code nodes are also added to nested graph
for edge in subgraph.edges():
if (isinstance(edge.src, nodes.CodeNode) and isinstance(edge.dst, nodes.CodeNode)):
subgraph_transients.add(edge.data.data)
# Collect data used in access nodes within subgraph (will be referenced in
# full upon nesting)
input_arrays = set()
output_arrays = {}
for node in subgraph.nodes():
if (isinstance(node, nodes.AccessNode) and node.data not in subgraph_transients):
if node.has_reads(state):
input_arrays.add(node.data)
if node.has_writes(state):
output_arrays[node.data] = state.in_edges(node)[0].data.wcr
# Create the nested SDFG
nsdfg = SDFG(name or 'nested_' + state.label)
# Transients are added to the nested graph as-is
for name in subgraph_transients:
nsdfg.add_datadesc(name, sdfg.arrays[name])
# Input/output data that are not source/sink nodes are added to the graph
# as non-transients
for name in (input_arrays | output_arrays.keys()):
datadesc = copy.deepcopy(sdfg.arrays[name])
datadesc.transient = False
nsdfg.add_datadesc(name, datadesc)
# Connected source/sink nodes outside subgraph become global data
# descriptors in nested SDFG
input_names = {}
output_names = {}
global_subsets: Dict[str, Tuple[str, Subset]] = {}
for edge in inputs:
if edge.data.data is None: # Skip edges with an empty memlet
continue
name = edge.data.data
if name not in global_subsets:
datadesc = copy.deepcopy(sdfg.arrays[edge.data.data])
datadesc.transient = False
if not full_data:
datadesc.shape = edge.data.subset.size()
new_name = nsdfg.add_datadesc(name, datadesc, find_new_name=True)
global_subsets[name] = (new_name, edge.data.subset)
else:
new_name, subset = global_subsets[name]
if not full_data:
new_subset = union(subset, edge.data.subset)
if new_subset is None:
new_subset = Range.from_array(sdfg.arrays[name])
global_subsets[name] = (new_name, new_subset)
nsdfg.arrays[new_name].shape = new_subset.size()
input_names[edge] = new_name
for edge in outputs:
if edge.data.data is None: # Skip edges with an empty memlet
continue
name = edge.data.data
if name not in global_subsets:
datadesc = copy.deepcopy(sdfg.arrays[edge.data.data])
datadesc.transient = False
if not full_data:
datadesc.shape = edge.data.subset.size()
new_name = nsdfg.add_datadesc(name, datadesc, find_new_name=True)
global_subsets[name] = (new_name, edge.data.subset)
else:
new_name, subset = global_subsets[name]
if not full_data:
new_subset = union(subset, edge.data.subset)
if new_subset is None:
new_subset = Range.from_array(sdfg.arrays[name])
global_subsets[name] = (new_name, new_subset)
nsdfg.arrays[new_name].shape = new_subset.size()
output_names[edge] = new_name
###################
# Add scope symbols to the nested SDFG
defined_vars = set(
symbolic.pystr_to_symbolic(s) for s in (state.symbols_defined_at(top_scopenode).keys()
| sdfg.symbols))
for v in defined_vars:
if v in sdfg.symbols:
sym = sdfg.symbols[v]
nsdfg.add_symbol(v, sym.dtype)
# Add constants to nested SDFG
for cstname, cstval in sdfg.constants.items():
nsdfg.add_constant(cstname, cstval)
# Create nested state
nstate = nsdfg.add_state()
# Add subgraph nodes and edges to nested state
nstate.add_nodes_from(subgraph.nodes())
for e in subgraph.edges():
nstate.add_edge(e.src, e.src_conn, e.dst, e.dst_conn, copy.deepcopy(e.data))
# Modify nested SDFG parents in subgraph
for node in subgraph.nodes():
if isinstance(node, nodes.NestedSDFG):
node.sdfg.parent = nstate
node.sdfg.parent_sdfg = nsdfg
node.sdfg.parent_nsdfg_node = node
# Add access nodes and edges as necessary
edges_to_offset = []
for edge, name in input_names.items():
node = nstate.add_read(name)
new_edge = copy.deepcopy(edge.data)
new_edge.data = name
edges_to_offset.append((edge, nstate.add_edge(node, None, edge.dst, edge.dst_conn, new_edge)))
for edge, name in output_names.items():
node = nstate.add_write(name)
new_edge = copy.deepcopy(edge.data)
new_edge.data = name
edges_to_offset.append((edge, nstate.add_edge(edge.src, edge.src_conn, node, None, new_edge)))
# Offset memlet paths inside nested SDFG according to subsets
for original_edge, new_edge in edges_to_offset:
for edge in nstate.memlet_tree(new_edge):
edge.data.data = new_edge.data.data
if not full_data:
edge.data.subset.offset(global_subsets[original_edge.data.data][1], True)
# Add nested SDFG node to the input state
nested_sdfg = state.add_nested_sdfg(nsdfg, None,
set(input_names.values()) | input_arrays,
set(output_names.values()) | output_arrays.keys())
# Reconnect memlets to nested SDFG
reconnected_in = set()
reconnected_out = set()
empty_input = None
empty_output = None
for edge in inputs:
if edge.data.data is None:
empty_input = edge
continue
name = input_names[edge]
if name in reconnected_in:
continue
if full_data:
data = Memlet.from_array(edge.data.data, sdfg.arrays[edge.data.data])
else:
data = copy.deepcopy(edge.data)
data.subset = global_subsets[edge.data.data][1]
state.add_edge(edge.src, edge.src_conn, nested_sdfg, name, data)
reconnected_in.add(name)
for edge in outputs:
if edge.data.data is None:
empty_output = edge
continue
name = output_names[edge]
if name in reconnected_out:
continue
if full_data:
data = Memlet.from_array(edge.data.data, sdfg.arrays[edge.data.data])
else:
data = copy.deepcopy(edge.data)
data.subset = global_subsets[edge.data.data][1]
data.wcr = edge.data.wcr
state.add_edge(nested_sdfg, name, edge.dst, edge.dst_conn, data)
reconnected_out.add(name)
# Connect access nodes to internal input/output data as necessary
entry = scope.entry
exit = scope.exit
for name in input_arrays:
node = state.add_read(name)
if entry is not None:
state.add_nedge(entry, node, Memlet())
state.add_edge(node, None, nested_sdfg, name, Memlet.from_array(name, sdfg.arrays[name]))
for name, wcr in output_arrays.items():
node = state.add_write(name)
if exit is not None:
state.add_nedge(node, exit, Memlet())
state.add_edge(nested_sdfg, name, node, None, Memlet(data=name, wcr=wcr))
# Graph was not reconnected, but needs to be
if state.in_degree(nested_sdfg) == 0 and empty_input is not None:
state.add_edge(empty_input.src, empty_input.src_conn, nested_sdfg, None, empty_input.data)
if state.out_degree(nested_sdfg) == 0 and empty_output is not None:
state.add_edge(nested_sdfg, None, empty_output.dst, empty_output.dst_conn, empty_output.data)
# Remove subgraph nodes from graph
state.remove_nodes_from(subgraph.nodes())
# Remove subgraph transients from top-level graph
for transient in subgraph_transients:
del sdfg.arrays[transient]
# Remove newly isolated nodes due to memlet consolidation
for edge in inputs:
if state.in_degree(edge.src) + state.out_degree(edge.src) == 0:
state.remove_node(edge.src)
for edge in outputs:
if state.in_degree(edge.dst) + state.out_degree(edge.dst) == 0:
state.remove_node(edge.dst)
return nested_sdfg
def _create_einsum_internal(sdfg: SDFG,
state: SDFGState,
einsum_string: str,
*arrays: str,
dtype: Optional[dtypes.typeclass] = None,
optimize: bool = False,
output: Optional[str] = None,
nodes: Optional[Dict[str, AccessNode]] = None,
init_output: bool = None):
# Infer shapes and strides of input/output arrays
einsum = EinsumParser(einsum_string)
if len(einsum.inputs) != len(arrays):
raise ValueError('Invalid number of arrays for einsum expression')
# Get shapes from arrays and verify dimensionality
chardict = {}
for inp, inpname in zip(einsum.inputs, arrays):
inparr = sdfg.arrays[inpname]
if len(inp) != len(inparr.shape):
raise ValueError('Dimensionality mismatch in input "%s"' % inpname)
for char, shp in zip(inp, inparr.shape):
if char in chardict and shp != chardict[char]:
raise ValueError('Dimension mismatch in einsum expression')
chardict[char] = shp
if optimize:
# Try to import opt_einsum
try:
import opt_einsum as oe
except (ModuleNotFoundError, NameError, ImportError):
raise ImportError('To optimize einsum expressions, please install '
'the "opt_einsum" package.')
for char, shp in chardict.items():
if symbolic.issymbolic(shp):
raise ValueError('Einsum optimization cannot be performed '
'on symbolically-sized array dimension "%s" '
'for subscript character "%s"' % (shp, char))
# Create optimal contraction path
# noinspection PyTypeChecker
_, path_info = oe.contract_path(
einsum_string, *oe.helpers.build_views(einsum_string, chardict))
input_nodes = nodes or {arr: state.add_read(arr) for arr in arrays}
result_node = None
# Follow path and create a chain of operation SDFG states
for pair, nonfree, expr, after, blas in path_info.contraction_list:
result, result_node = _create_einsum_internal(sdfg,
state,
expr,
arrays[pair[0]],
arrays[pair[1]],
dtype=dtype,
optimize=False,
output=None,
nodes=input_nodes)
arrays = ([a for i, a in enumerate(arrays) if i not in pair] +
[result])
input_nodes[result] = result_node
return arrays[0], result_node
# END of einsum optimization
input_nodes = nodes or {arr: state.add_read(arr) for arr in arrays}
# Get output shape from chardict, or [1] for a scalar output
output_shape = list(map(lambda k: chardict[k], einsum.output)) or [1]
output_index = ','.join(o for o in einsum.output) or '0'
if output is None:
dtype = dtype or sdfg.arrays[arrays[0]].dtype
output, odesc = sdfg.add_temp_transient(output_shape, dtype)
to_init = True
else:
odesc = sdfg.arrays[output]
dtype = dtype or odesc.dtype
to_init = init_output or True
is_conflicted = not all(
all(indim in einsum.output for indim in inp) for inp in einsum.inputs)
if not is_conflicted and init_output is None:
to_init = False
if not einsum.is_bmm():
# Fall back to "pure" SDFG einsum with conflict resolution
c = state.add_write(output)
# Add state before this one to initialize the output value
if to_init:
init_state = sdfg.add_state_before(state)
if len(einsum.output) > 0:
init_state.add_mapped_tasklet(
'einsum_reset',
{k: '0:%s' % chardict[k]
for k in einsum.output}, {},
'out_%s = 0' % output,
{'out_%s' % output: Memlet.simple(output, output_index)},
external_edges=True)
else: # Scalar output
t = init_state.add_tasklet('einsum_reset', set(),
{'out_%s' % output},
'out_%s = 0' % output)
onode = init_state.add_write(output)
init_state.add_edge(t, 'out_%s' % output, onode, None,
Memlet.simple(output, '0'))
wcr = 'lambda a,b: a+b' if is_conflicted else None
# Pure einsum map
state.add_mapped_tasklet(
'einsum', {k: '0:%s' % v
for k, v in chardict.items()}, {
'inp_%s' % arr: Memlet.simple(arr, ','.join(inp))
for inp, arr in zip(einsum.inputs, arrays)
},
'out_%s = %s' % (output, ' * '.join('inp_%s' % arr
for arr in arrays)),
{
'out_%s' % output: Memlet.simple(
output, output_index, wcr_str=wcr)
},
input_nodes=input_nodes,
output_nodes={output: c},
external_edges=True)
else:
# Represent einsum as a GEMM or batched GEMM (using library nodes)
a_shape = sdfg.arrays[arrays[0]].shape
b_shape = sdfg.arrays[arrays[1]].shape
c_shape = output_shape
a = input_nodes[arrays[0]]
b = input_nodes[arrays[1]]
c = state.add_write(output)
# Compute GEMM dimensions and strides
strides = dict(
BATCH=prod([c_shape[dim] for dim in einsum.c_batch]),
M=prod([a_shape[dim] for dim in einsum.a_only]),
K=prod([a_shape[dim] for dim in einsum.a_sum]),
N=prod([b_shape[dim] for dim in einsum.b_only]),
sAM=prod(a_shape[einsum.a_only[-1] + 1:]) if einsum.a_only else 1,
sAK=prod(a_shape[einsum.a_sum[-1] + 1:]) if einsum.a_sum else 1,
sAB=prod(a_shape[einsum.a_batch[-1] +
1:]) if einsum.a_batch else 1,
sBK=prod(b_shape[einsum.b_sum[-1] + 1:]) if einsum.b_sum else 1,
sBN=prod(b_shape[einsum.b_only[-1] + 1:]) if einsum.b_only else 1,
sBB=prod(b_shape[einsum.b_batch[-1] +
1:]) if einsum.b_batch else 1,
sCM=prod(c_shape[einsum.c_a_only[-1] +
1:]) if einsum.c_a_only else 1,
sCN=prod(c_shape[einsum.c_b_only[-1] +
1:]) if einsum.c_b_only else 1,
sCB=prod(c_shape[einsum.c_batch[-1] +
1:]) if einsum.c_batch else 1)
# Complement strides to make matrices as necessary
if len(a_shape) == 1 and len(einsum.a_sum) == 1:
strides['sAK'] = 1
strides['sAB'] = strides['sAM'] = strides['K']
if len(b_shape) == 1 and len(einsum.b_sum) == 1:
strides['sBN'] = 1
strides['sBK'] = 1
strides['sBB'] = strides['K']
if len(c_shape) == 1 and len(einsum.a_sum) == len(einsum.b_sum):
strides['sCN'] = 1
strides['sCB'] = strides['sCM'] = strides['N']
# Create nested SDFG for GEMM
nsdfg = create_batch_gemm_sdfg(dtype, strides)
nsdfg_node = state.add_nested_sdfg(nsdfg, None, {'X', 'Y'}, {'Z'},
strides)
state.add_edge(a, None, nsdfg_node, 'X',
Memlet.from_array(a.data, a.desc(sdfg)))
state.add_edge(b, None, nsdfg_node, 'Y',
Memlet.from_array(b.data, b.desc(sdfg)))
state.add_edge(nsdfg_node, 'Z', c, None,
Memlet.from_array(c.data, c.desc(sdfg)))
return output, c
