def can_be_applied(self,
graph: SDFGState,
expr_index: int,
sdfg: SDFG,
permissive: bool = False) -> bool:
access = self.access
# Make sure the access node is only accessed once (read or write),
# and not at the same time
if graph.out_degree(access) > 0 and graph.in_degree(access) > 0:
return False
# If already a stream, skip
if isinstance(sdfg.arrays[access.data], data.Stream):
return False
# If does not exist on off-chip memory, skip
if sdfg.arrays[access.data].storage not in [
dtypes.StorageType.CPU_Heap, dtypes.StorageType.CPU_Pinned,
dtypes.StorageType.GPU_Global, dtypes.StorageType.FPGA_Global
]:
return False
# Only free nodes are allowed (search up the SDFG tree)
curstate = graph
node = access
while curstate is not None:
if curstate.entry_node(node) is not None:
return False
if curstate.parent.parent_nsdfg_node is None:
break
node = curstate.parent.parent_nsdfg_node
curstate = curstate.parent.parent
# Only one memlet path is allowed per outgoing/incoming edge
edges = (graph.out_edges(access)
if expr_index == 0 else graph.in_edges(access))
for edge in edges:
mpath = graph.memlet_path(edge)
if len(mpath) != len(list(graph.memlet_tree(edge))):
return False
# The innermost end of the path must have a clearly defined memory
# access pattern
innermost_edge = mpath[-1] if expr_index == 0 else mpath[0]
if (innermost_edge.data.subset.num_elements() != 1
or innermost_edge.data.dynamic
or innermost_edge.data.volume != 1):
return False
# Check if any of the maps has a dynamic range
# These cases can potentially work but some nodes (and perhaps
# tasklets) need to be replicated, which are difficult to track.
for pe in mpath:
node = pe.dst if expr_index == 0 else graph.entry_node(pe.src)
if isinstance(
node,
nodes.MapEntry) and sdutil.has_dynamic_map_inputs(
graph, node):
return False
# If already applied on this memlet and this is the I/O component, skip
if expr_index == 0:
other_node = self.entry
else:
other_node = self.exit
other_node = graph.entry_node(other_node)
if other_node.label.startswith('__s'):
return False
## Check Memory Buffering Properties
if self.use_memory_buffering:
access = self.access
desc = sdfg.arrays[access.data]
# Array has to be global array
if desc.storage != dtypes.StorageType.FPGA_Global:
return False
# Type has to divide target bytes
if self.memory_buffering_target_bytes % desc.dtype.bytes != 0:
return False
# Target bytes has to be >= size of data type
if self.memory_buffering_target_bytes < desc.dtype.bytes:
return False
strides = list(desc.strides)
# Last stride has to be one
if strides[-1] != 1:
return False
vector_size = int(self.memory_buffering_target_bytes /
desc.dtype.bytes)
strides.pop() # Remove last element since we already checked it
# Other strides have to be divisible by vector size
for stride in strides:
if is_int(stride) and stride % vector_size != 0:
return False
# Check if map has the right access pattern
# Stride 1 access by innermost loop, innermost loop counter has to be divisible by vector size
# Same code as in apply
state = sdfg.node(self.state_id)
dnode: nodes.AccessNode = self.access
if self.expr_index == 0:
edges = state.out_edges(dnode)
else:
edges = state.in_edges(dnode)
mapping: Dict[
Tuple[subsets.Range],
List[gr.MultiConnectorEdge[mm.Memlet]]] = defaultdict(list)
ranges = {}
for edge in edges:
mpath = state.memlet_path(edge)
ranges[edge] = _collect_map_ranges(state, mpath)
mapping[tuple(r[1] for r in ranges[edge])].append(edge)
for edges_with_same_range in mapping.values():
for edge in edges_with_same_range:
# Get memlet path and innermost edge
mpath = state.memlet_path(edge)
innermost_edge = copy.deepcopy(
mpath[-1] if self.expr_index == 0 else mpath[0])
edge_subset = [
a_tuple[0]
for a_tuple in list(innermost_edge.data.subset)
]
if self.expr_index == 0:
map_subset = innermost_edge.src.map.params.copy()
ranges = list(innermost_edge.src.map.range)
else:
map_subset = innermost_edge.dst.map.params.copy()
ranges = list(innermost_edge.dst.map.range)
# Check is correct access pattern
# Correct ranges in map
if is_int(ranges[-1]
[1]) and (ranges[-1][1] + 1) % vector_size != 0:
return False
if ranges[-1][2] != 1:
return False
# Correct access in array
if isinstance(edge_subset[-1], symbol) and str(
edge_subset[-1]) == map_subset[-1]:
pass
elif isinstance(edge_subset[-1], sympy.core.add.Add):
counter: int = 0
for arg in edge_subset[-1].args:
if isinstance(
arg,
symbol) and str(arg) == map_subset[-1]:
counter += 1
if counter != 1:
return False
else:
return False
return True
def can_be_applied(graph: SDFGState,
candidate: Dict[xf.PatternNode, int],
expr_index: int,
sdfg: SDFG,
strict: bool = False) -> bool:
access = graph.node(candidate[StreamingComposition.access])
# Make sure the access node is only accessed once (read or write),
# and not at the same time
if graph.in_degree(access) > 1 or graph.out_degree(access) > 1:
return False
# If already a stream, skip
if isinstance(sdfg.arrays[access.data], data.Stream):
return False
# Only free nodes are allowed (search up the SDFG tree)
curstate = graph
node = access
while curstate is not None:
if curstate.entry_node(node) is not None:
return False
if curstate.parent.parent_nsdfg_node is None:
break
node = curstate.parent.parent_nsdfg_node
curstate = curstate.parent.parent
# Array must not be used anywhere else in the state
if any(n is not access and n.data == access.data
for n in graph.data_nodes()):
return False
# Only one memlet path on each direction is allowed
# TODO: Relax so that repeated application of
# transformation would yield additional streams
first_edge = graph.in_edges(access)[0]
second_edge = graph.out_edges(access)[0]
first_mpath = graph.memlet_path(first_edge)
second_mpath = graph.memlet_path(second_edge)
if len(first_mpath) != len(list(graph.memlet_tree(first_edge))):
return False
if len(second_mpath) != len(list(graph.memlet_tree(second_edge))):
return False
# The innermost ends of the paths must have a clearly defined memory
# access pattern and no WCR
first_iedge = first_mpath[0]
second_iedge = second_mpath[-1]
if first_iedge.data.subset.num_elements() != 1:
return False
if first_iedge.data.volume != 1:
return False
if first_iedge.data.wcr is not None:
return False
if second_iedge.data.subset.num_elements() != 1:
return False
if second_iedge.data.volume != 1:
return False
##################################################################
# The memory access pattern must be exactly the same
# Collect all maps and ranges
ranges_first = _collect_map_ranges(graph, first_mpath)
ranges_second = _collect_map_ranges(graph, second_mpath)
# Check map ranges
for (_, frng), (_, srng) in zip(ranges_first, ranges_second):
if frng != srng:
return False
# Check memlets for equivalence
if len(first_iedge.data.subset) != len(second_iedge.data.subset):
return False
if not _do_memlets_correspond(first_iedge.data, second_iedge.data,
ranges_first, ranges_second):
return False
return True
def can_be_applied(self,
graph: SDFGState,
expr_index,
sdfg,
permissive=False):
nsdfg = self.nsdfg
# Must be a free nested SDFG
if graph.entry_node(nsdfg) is not None:
return False
# Must have two states with an empty source state.
# Otherwise structured control flow (loop init states, for example)
# may be broken.
if not permissive:
if nsdfg.sdfg.number_of_nodes() != 2:
return False
if nsdfg.sdfg.start_state.number_of_nodes() != 0:
return False
# Must have at least two states with a hoistable source state
if nsdfg.sdfg.number_of_nodes() < 2:
return False
# Source state must not lead to more than one state or be conditional
source_state = nsdfg.sdfg.start_state
if nsdfg.sdfg.out_degree(source_state) != 1:
return False
nisedge = nsdfg.sdfg.out_edges(source_state)[0]
if not nisedge.data.is_unconditional():
return False
# Keep all data descriptors to check for potential issues
data_to_check: Set[str] = set()
# Add data descriptors from interstate edge
syms = nisedge.data.free_symbols
for sym in syms:
sym = str(sym)
if sym in nsdfg.sdfg.arrays:
if nsdfg.sdfg.arrays[sym].transient: # Cannot keep transient
return False
data_to_check.add(sym)
# Add data descriptors from access nodes
for dnode in source_state.data_nodes():
data_to_check.add(dnode.data)
desc = nsdfg.sdfg.arrays[dnode.data]
# Cannot hoist state with transient
if not isinstance(desc, dt.View) and desc.transient:
return False
# Nested SDFG surrounding edges must contain all of the array
# TODO(later): Allow this case (with offsetting)
outer_data_to_check: Set[str] = set()
for e in graph.in_edges(nsdfg):
if e.dst_conn in data_to_check:
outer_data_to_check.add(e.data.data)
if any(me != 0 for me in e.data.subset.min_element()):
return False
for e in graph.out_edges(nsdfg):
if e.src_conn in data_to_check:
outer_data_to_check.add(e.data.data)
if any(me != 0 for me in e.data.subset.min_element()):
return False
# Data validity checks for descriptors in data_to_check:
# 1. Path to nested SDFG must not go through descriptors,
# 2. No other connected components can use descriptors.
for dnode in graph.data_nodes():
if dnode.data in outer_data_to_check:
if nx.has_path(graph._nx, nsdfg, dnode):
# OK, has path from nsdfg to access node
continue
if dnode in graph.predecessors(nsdfg):
# OK, a direct edge to nsdfg
continue
if nx.has_path(graph._nx, dnode, nsdfg):
# NOT OK, some path goes through access node to SDFG,
# so state cannot safely be hoisted
return False
# NOT OK, access node used independently from nsdfg
return False
return True
def expansion(node: 'Reduce', state: SDFGState, sdfg: SDFG):
from dace.codegen.prettycode import CodeIOStream
from dace.codegen.targets.cpp import unparse_cr_split, cpp_array_expr
node.validate(sdfg, state)
input_edge: graph.MultiConnectorEdge = state.in_edges(node)[0]
output_edge: graph.MultiConnectorEdge = state.out_edges(node)[0]
input_dims = len(input_edge.data.subset)
input_data = sdfg.arrays[input_edge.data.data]
output_data = sdfg.arrays[output_edge.data.data]
# Setup all locations in which code will be written
cuda_globalcode = CodeIOStream()
localcode = CodeIOStream()
# Try to autodetect reduction type
redtype = detect_reduction_type(node.wcr)
node_id = state.node_id(node)
state_id = sdfg.node_id(state)
idstr = '{sdfg}_{state}_{node}'.format(sdfg=sdfg.name,
state=state_id,
node=node_id)
# Obtain some SDFG-related information
input_memlet = input_edge.data
output_memlet = output_edge.data
if node.out_connectors:
dtype = next(node.out_connectors.values())
else:
dtype = sdfg.arrays[output_memlet.data].dtype
output_type = dtype.ctype
if node.identity is None:
raise ValueError('For device reduce nodes, initial value must be '
'specified')
# Create a functor or use an existing one for reduction
if redtype == dtypes.ReductionType.Custom:
body, [arg1, arg2] = unparse_cr_split(sdfg, node.wcr)
cuda_globalcode.write(
"""
struct __reduce_{id} {{
template <typename T>
DACE_HDFI T operator()(const T &{arg1}, const T &{arg2}) const {{
{contents}
}}
}};""".format(id=idstr, arg1=arg1, arg2=arg2, contents=body), sdfg,
state_id, node_id)
reduce_op = ', __reduce_' + idstr + '(), ' + symstr(node.identity)
elif redtype in ExpandReduceCUDADevice._SPECIAL_RTYPES:
reduce_op = ''
else:
credtype = 'dace::ReductionType::' + str(
redtype)[str(redtype).find('.') + 1:]
reduce_op = ((', dace::_wcr_fixed<%s, %s>()' %
(credtype, output_type)) + ', ' +
symstr(node.identity))
# Try to obtain the number of threads in the block, or use the default
# configuration
block_threads = devicelevel_block_size(sdfg, state, node)
if block_threads is not None:
block_threads = functools.reduce(lambda a, b: a * b, block_threads,
1)
# Checks
if block_threads is None:
raise ValueError('Block-wide GPU reduction must occur within'
' a GPU kernel')
if issymbolic(block_threads, sdfg.constants):
raise ValueError('Block size has to be constant for block-wide '
'reduction (got %s)' % str(block_threads))
if (node.axes is not None and len(node.axes) < input_dims):
raise ValueError(
'Only full reduction is supported for block-wide reduce,'
' please use the pure expansion')
if (input_data.storage != dtypes.StorageType.Register
or output_data.storage != dtypes.StorageType.Register):
raise ValueError(
'Block-wise reduction only supports GPU register inputs '
'and outputs')
if redtype in ExpandReduceCUDABlock._SPECIAL_RTYPES:
raise ValueError('%s block reduction not supported' % redtype)
credtype = 'dace::ReductionType::' + str(
redtype)[str(redtype).find('.') + 1:]
if redtype == dtypes.ReductionType.Custom:
redop = '__reduce_%s()' % idstr
else:
redop = 'dace::_wcr_fixed<%s, %s>()' % (credtype, output_type)
# Allocate shared memory for block reduce
localcode.write("""
typedef cub::BlockReduce<{type}, {numthreads}> BlockReduce_{id};
__shared__ typename BlockReduce_{id}::TempStorage temp_storage_{id};
""".format(id=idstr,
type=output_data.dtype.ctype,
numthreads=block_threads))
input = (input_memlet.data + ' + ' +
cpp_array_expr(sdfg, input_memlet, with_brackets=False))
output = cpp_array_expr(sdfg, output_memlet)
localcode.write("""
{output} = BlockReduce_{id}(temp_storage_{id}).Reduce({input}, {redop});
""".format(id=idstr,
redop=redop,
input=input_memlet.data,
output=output))
# Make tasklet
tnode = dace.nodes.Tasklet('reduce',
{'_in': dace.pointer(input_data.dtype)},
{'_out': dace.pointer(output_data.dtype)},
localcode.getvalue(),
language=dace.Language.CPP)
# Add the rest of the code
sdfg.append_global_code(cuda_globalcode.getvalue(), 'cuda')
# Rename outer connectors and add to node
input_edge._dst_conn = '_in'
output_edge._src_conn = '_out'
node.add_in_connector('_in')
node.add_out_connector('_out')
return tnode
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
def expansion(node: 'Reduce', state: SDFGState, sdfg: SDFG):
node.validate(sdfg, state)
inedge: graph.MultiConnectorEdge = state.in_edges(node)[0]
outedge: graph.MultiConnectorEdge = state.out_edges(node)[0]
insubset = dcpy(inedge.data.subset)
isqdim = insubset.squeeze()
outsubset = dcpy(outedge.data.subset)
osqdim = outsubset.squeeze()
input_dims = len(insubset)
output_dims = len(outsubset)
input_data = sdfg.arrays[inedge.data.data]
output_data = sdfg.arrays[outedge.data.data]
if len(osqdim) == 0: # Fix for scalars
osqdim = [0]
# Standardize axes
axes = node.axes if node.axes else [i for i in range(input_dims)]
assert node.identity is not None
assert len(axes) == 1
# Create nested SDFG
nsdfg = SDFG('reduce')
nsdfg.add_array('_in',
insubset.size(),
input_data.dtype,
strides=[
s for i, s in enumerate(input_data.strides)
if i in isqdim
],
storage=input_data.storage)
nsdfg.add_array('_out',
outsubset.size(),
output_data.dtype,
strides=[
s for i, s in enumerate(output_data.strides)
if i in osqdim
],
storage=output_data.storage)
nsdfg.add_transient('acc', [1], nsdfg.arrays['_in'].dtype,
dtypes.StorageType.Register)
nstate = nsdfg.add_state()
axis = axes[0]
inp = dcpy(nsdfg.arrays['_in'])
out = dcpy(nsdfg.arrays['_out'])
# Interleave input and output axes to match input memlet
ictr, octr = 0, 0
input_subset = []
for i in isqdim:
if i in axes:
input_subset.append('_i%d' % ictr)
ictr += 1
else:
input_subset.append('_o%d' % octr)
octr += 1
ome, omx = nstate.add_map(
'reduce_output', {
'_o%d' % i: '0:%s' % symstr(sz)
for i, sz in enumerate(outsubset.size())
})
outm = dace.Memlet.simple(
'_out', ','.join(['_o%d' % i for i in range(output_dims)]))
#wcr_str=node.wcr)
inmm = dace.Memlet.simple('_in', ','.join(input_subset))
idt = nstate.add_tasklet('reset', {}, {'o'}, f'o = {node.identity}')
nstate.add_edge(ome, None, idt, None, dace.Memlet())
accread = nstate.add_access('acc')
accwrite = nstate.add_access('acc')
nstate.add_edge(idt, 'o', accread, None, dace.Memlet('acc'))
# Add inner map, which corresponds to the range to reduce, containing
# an identity tasklet
ime, imx = nstate.add_map('reduce_values', {
'_i%d' % i: '0:%s' % symstr(insubset.size()[isqdim.index(axis)])
for i, axis in enumerate(sorted(axes))
},
schedule=dtypes.ScheduleType.Sequential)
# Add identity tasklet for reduction
t = nstate.add_tasklet('identity', {'a', 'b'}, {'o'}, 'o = a + b')
# Connect everything
r = nstate.add_read('_in')
w = nstate.add_write('_out')
nstate.add_memlet_path(r, ome, ime, t, dst_conn='b', memlet=inmm)
nstate.add_memlet_path(accread,
ime,
t,
dst_conn='a',
memlet=dace.Memlet('acc[0]'))
nstate.add_memlet_path(t,
imx,
accwrite,
src_conn='o',
memlet=dace.Memlet('acc[0]'))
nstate.add_memlet_path(accwrite, omx, w, memlet=outm)
# Rename outer connectors and add to node
inedge._dst_conn = '_in'
outedge._src_conn = '_out'
node.add_in_connector('_in')
node.add_out_connector('_out')
from dace.transformation import dataflow
nsdfg.apply_transformations_repeated(dataflow.MapCollapse)
return nsdfg
def expansion(node: 'Reduce', state: SDFGState, sdfg: SDFG):
node.validate(sdfg, state)
inedge: graph.MultiConnectorEdge = state.in_edges(node)[0]
outedge: graph.MultiConnectorEdge = state.out_edges(node)[0]
insubset = dcpy(inedge.data.subset)
isqdim = insubset.squeeze()
outsubset = dcpy(outedge.data.subset)
osqdim = outsubset.squeeze()
input_dims = len(insubset)
output_dims = len(outsubset)
input_data = sdfg.arrays[inedge.data.data]
output_data = sdfg.arrays[outedge.data.data]
if len(osqdim) == 0: # Fix for scalars
osqdim = [0]
# Standardize and squeeze axes
axes = node.axes if node.axes else [
i for i in range(len(inedge.data.subset))
]
axes = [axis for axis in axes if axis in isqdim]
# Create nested SDFG
nsdfg = SDFG('reduce')
nsdfg.add_array('_in',
insubset.size(),
input_data.dtype,
strides=[
s for i, s in enumerate(input_data.strides)
if i in isqdim
],
storage=input_data.storage)
nsdfg.add_array('_out',
outsubset.size(),
output_data.dtype,
strides=[
s for i, s in enumerate(output_data.strides)
if i in osqdim
],
storage=output_data.storage)
# If identity is defined, add an initialization state
if node.identity is not None:
init_state = nsdfg.add_state()
nstate = nsdfg.add_state()
nsdfg.add_edge(init_state, nstate, dace.InterstateEdge())
# Add initialization as a map
init_state.add_mapped_tasklet(
'reduce_init', {
'_o%d' % i: '0:%s' % symstr(d)
for i, d in enumerate(outedge.data.subset.size())
}, {},
'out = %s' % node.identity, {
'out':
dace.Memlet.simple(
'_out', ','.join(
['_o%d' % i for i in range(output_dims)]))
},
external_edges=True)
else:
nstate = nsdfg.add_state()
# END OF INIT
# (If axes != all) Add outer map, which corresponds to the output range
if len(axes) != input_dims:
# Interleave input and output axes to match input memlet
ictr, octr = 0, 0
input_subset = []
for i in isqdim:
if i in axes:
input_subset.append('_i%d' % ictr)
ictr += 1
else:
input_subset.append('_o%d' % octr)
octr += 1
ome, omx = nstate.add_map(
'reduce_output', {
'_o%d' % i: '0:%s' % symstr(sz)
for i, sz in enumerate(outsubset.size())
})
outm = dace.Memlet.simple(
'_out',
','.join(['_o%d' % i for i in range(output_dims)]),
wcr_str=node.wcr)
inmm = dace.Memlet.simple('_in', ','.join(input_subset))
else:
ome, omx = None, None
outm = dace.Memlet.simple('_out', '0', wcr_str=node.wcr)
inmm = dace.Memlet.simple(
'_in', ','.join(['_i%d' % i for i in range(len(axes))]))
# Add inner map, which corresponds to the range to reduce, containing
# an identity tasklet
ime, imx = nstate.add_map(
'reduce_values', {
'_i%d' % i:
'0:%s' % symstr(insubset.size()[isqdim.index(axis)])
for i, axis in enumerate(sorted(axes))
})
# Add identity tasklet for reduction
t = nstate.add_tasklet('identity', {'inp'}, {'out'}, 'out = inp')
# Connect everything
r = nstate.add_read('_in')
w = nstate.add_read('_out')
if ome:
nstate.add_memlet_path(r, ome, ime, t, dst_conn='inp', memlet=inmm)
nstate.add_memlet_path(t, imx, omx, w, src_conn='out', memlet=outm)
else:
nstate.add_memlet_path(r, ime, t, dst_conn='inp', memlet=inmm)
nstate.add_memlet_path(t, imx, w, src_conn='out', memlet=outm)
# Rename outer connectors and add to node
inedge._dst_conn = '_in'
outedge._src_conn = '_out'
node.add_in_connector('_in')
node.add_out_connector('_out')
from dace.transformation import dataflow
nsdfg.apply_transformations_repeated(dataflow.MapCollapse)
return nsdfg
def _bcgather(pv: 'ProgramVisitor', sdfg: SDFG, state: SDFGState,
in_buffer: str, out_buffer: str,
block_sizes: Union[str, Sequence[Union[sp.Expr, Number]]]):
from dace.libraries.pblas.nodes.pgeadd import BlockCyclicGather
libnode = BlockCyclicGather('_BCGather_')
inbuf_range = None
if isinstance(in_buffer, tuple):
inbuf_name, inbuf_range = in_buffer
else:
inbuf_name = in_buffer
in_desc = sdfg.arrays[inbuf_name]
inbuf_node = state.add_read(inbuf_name)
bsizes_range = None
if isinstance(block_sizes, (list, tuple)):
if isinstance(block_sizes[0], str):
bsizes_name, bsizes_range = block_sizes
bsizes_desc = sdfg.arrays[bsizes_name]
bsizes_node = state.add_read(bsizes_name)
else:
bsizes_name, bsizes_desc = sdfg.add_temp_transient(
(len(block_sizes), ), dtype=dace.int32)
bsizes_node = state.add_access(bsizes_name)
bsizes_tasklet = state.add_tasklet(
'_set_bsizes_', {}, {'__out'}, ";".join([
"__out[{}] = {}".format(i, sz)
for i, sz in enumerate(block_sizes)
]))
state.add_edge(bsizes_tasklet, '__out', bsizes_node, None,
Memlet.from_array(bsizes_name, bsizes_desc))
else:
bsizes_name = block_sizes
bsizes_desc = sdfg.arrays[bsizes_name]
bsizes_node = state.add_read(bsizes_name)
outbuf_range = None
if isinstance(out_buffer, tuple):
outbuf_name, outbuf_range = out_buffer
else:
outbuf_name = out_buffer
out_desc = sdfg.arrays[outbuf_name]
outbuf_node = state.add_write(outbuf_name)
if inbuf_range:
inbuf_mem = Memlet.simple(inbuf_name, inbuf_range)
else:
inbuf_mem = Memlet.from_array(inbuf_name, in_desc)
if bsizes_range:
bsizes_mem = Memlet.simple(bsizes_name, bsizes_range)
else:
bsizes_mem = Memlet.from_array(bsizes_name, bsizes_desc)
if outbuf_range:
outbuf_mem = Memlet.simple(outbuf_name, outbuf_range)
else:
outbuf_mem = Memlet.from_array(outbuf_name, out_desc)
state.add_edge(inbuf_node, None, libnode, '_inbuffer', inbuf_mem)
state.add_edge(bsizes_node, None, libnode, '_block_sizes', bsizes_mem)
state.add_edge(libnode, '_outbuffer', outbuf_node, None, outbuf_mem)
return None
def _distr_matmult(pv: 'ProgramVisitor',
sdfg: SDFG,
state: SDFGState,
opa: str,
opb: str,
shape: Sequence[Union[sp.Expr, Number]],
a_block_sizes: Union[str, Sequence[Union[sp.Expr,
Number]]] = None,
b_block_sizes: Union[str, Sequence[Union[sp.Expr,
Number]]] = None,
c_block_sizes: Union[str, Sequence[Union[sp.Expr,
Number]]] = None):
arra = sdfg.arrays[opa]
arrb = sdfg.arrays[opb]
if len(shape) == 3:
gm, gn, gk = shape
else:
gm, gn = shape
a_block_sizes = a_block_sizes or arra.shape
if len(a_block_sizes) < 2:
a_block_sizes = (a_block_sizes[0], 1)
b_block_sizes = b_block_sizes or arrb.shape
if len(b_block_sizes) < 2:
b_block_sizes = (b_block_sizes[0], 1)
if len(arra.shape) == 1 and len(arrb.shape) == 2:
a_block_sizes, b_block_sizes = b_block_sizes, a_block_sizes
a_bsizes_range = None
if isinstance(a_block_sizes, (list, tuple)):
if isinstance(a_block_sizes[0], str):
a_bsizes_name, a_bsizes_range = a_block_sizes
a_bsizes_desc = sdfg.arrays[a_bsizes_name]
a_bsizes_node = state.add_read(a_bsizes_name)
else:
a_bsizes_name, a_bsizes_desc = sdfg.add_temp_transient(
(len(a_block_sizes), ), dtype=dace.int32)
a_bsizes_node = state.add_access(a_bsizes_name)
a_bsizes_tasklet = state.add_tasklet(
'_set_a_bsizes_', {}, {'__out'}, ";".join([
"__out[{}] = {}".format(i, sz)
for i, sz in enumerate(a_block_sizes)
]))
state.add_edge(a_bsizes_tasklet, '__out', a_bsizes_node, None,
Memlet.from_array(a_bsizes_name, a_bsizes_desc))
else:
a_bsizes_name = a_block_sizes
a_bsizes_desc = sdfg.arrays[a_bsizes_name]
a_bsizes_node = state.add_read(a_bsizes_name)
b_bsizes_range = None
if isinstance(a_block_sizes, (list, tuple)):
if isinstance(a_block_sizes[0], str):
b_bsizes_name, b_sizes_range = b_block_sizes
b_bsizes_desc = sdfg.arrays[b_bsizes_name]
b_bsizes_node = state.add_read(b_bsizes_name)
else:
b_bsizes_name, b_bsizes_desc = sdfg.add_temp_transient(
(len(b_block_sizes), ), dtype=dace.int32)
b_bsizes_node = state.add_access(b_bsizes_name)
b_bsizes_tasklet = state.add_tasklet(
'_set_b_sizes_', {}, {'__out'}, ";".join([
"__out[{}] = {}".format(i, sz)
for i, sz in enumerate(b_block_sizes)
]))
state.add_edge(b_bsizes_tasklet, '__out', b_bsizes_node, None,
Memlet.from_array(b_bsizes_name, b_bsizes_desc))
else:
b_bsizes_name = b_block_sizes
b_bsizes_desc = sdfg.arrays[b_bsizes_name]
b_bsizes_node = state.add_read(b_bsizes_name)
if len(arra.shape) == 2 and len(arrb.shape) == 2:
# Gemm
from dace.libraries.pblas.nodes.pgemm import Pgemm
tasklet = Pgemm("__DistrMatMult__", gm, gn, gk)
m = arra.shape[0]
n = arrb.shape[-1]
out = sdfg.add_temp_transient((m, n), dtype=arra.dtype)
elif len(arra.shape) == 2 and len(arrb.shape) == 1:
# Gemv
from dace.libraries.pblas.nodes.pgemv import Pgemv
tasklet = Pgemv("__DistrMatVecMult__", m=gm, n=gn)
if c_block_sizes:
m = c_block_sizes[0]
else:
m = arra.shape[0]
out = sdfg.add_temp_transient((m, ), dtype=arra.dtype)
elif len(arra.shape) == 1 and len(arrb.shape) == 2:
# Gemv transposed
# Swap a and b
opa, opb = opb, opa
arra, arrb = arrb, arra
from dace.libraries.pblas.nodes.pgemv import Pgemv
tasklet = Pgemv("__DistrMatVecMult__", transa='T', m=gm, n=gn)
if c_block_sizes:
n = c_block_sizes[0]
else:
n = arra.shape[1]
out = sdfg.add_temp_transient((n, ), dtype=arra.dtype)
anode = state.add_read(opa)
bnode = state.add_read(opb)
cnode = state.add_write(out[0])
if a_bsizes_range:
a_bsizes_mem = Memlet.simple(a_bsizes_name, a_bsizes_range)
else:
a_bsizes_mem = Memlet.from_array(a_bsizes_name, a_bsizes_desc)
if b_bsizes_range:
b_bsizes_mem = Memlet.simple(b_bsizes_name, b_bsizes_range)
else:
b_bsizes_mem = Memlet.from_array(b_bsizes_name, b_bsizes_desc)
state.add_edge(anode, None, tasklet, '_a', Memlet.from_array(opa, arra))
state.add_edge(bnode, None, tasklet, '_b', Memlet.from_array(opb, arrb))
state.add_edge(a_bsizes_node, None, tasklet, '_a_block_sizes',
a_bsizes_mem)
state.add_edge(b_bsizes_node, None, tasklet, '_b_block_sizes',
b_bsizes_mem)
state.add_edge(tasklet, '_c', cnode, None, Memlet.from_array(*out))
return out[0]
def _isend(pv: 'ProgramVisitor', sdfg: SDFG, state: SDFGState, buffer: str,
dst: Union[str, sp.Expr, Number], tag: Union[str, sp.Expr,
Number], request: str):
from dace.libraries.mpi.nodes.isend import Isend
libnode = Isend('_Isend_')
buf_range = None
if isinstance(buffer, tuple):
buf_name, buf_range = buffer
else:
buf_name = buffer
desc = sdfg.arrays[buf_name]
buf_node = state.add_read(buf_name)
req_range = None
if isinstance(request, tuple):
req_name, req_range = request
else:
req_name = request
req_desc = sdfg.arrays[req_name]
req_node = state.add_write(req_name)
iconn = libnode.in_connectors
iconn = {
c: (dtypes.pointer(desc.dtype) if c == '_buffer' else t)
for c, t in iconn.items()
}
libnode.in_connectors = iconn
oconn = libnode.out_connectors
oconn = {
c: (dtypes.pointer(req_desc.dtype) if c == '_request' else t)
for c, t in oconn.items()
}
libnode.out_connectors = oconn
dst_range = None
if isinstance(dst, tuple):
dst_name, dst_range = dst
dst_node = state.add_read(dst_name)
elif isinstance(dst, str) and dst in sdfg.arrays.keys():
dst_name = dst
dst_node = state.add_read(dst_name)
else:
storage = desc.storage
dst_name = _define_local_scalar(pv, sdfg, state, dace.int32, storage)
dst_node = state.add_access(dst_name)
dst_tasklet = state.add_tasklet('_set_dst_', {}, {'__out'},
'__out = {}'.format(dst))
state.add_edge(dst_tasklet, '__out', dst_node, None,
Memlet.simple(dst_name, '0'))
tag_range = None
if isinstance(tag, tuple):
tag_name, tag_range = tag
tag_node = state.add_read(tag_name)
if isinstance(tag, str) and tag in sdfg.arrays.keys():
tag_name = tag
tag_node = state.add_read(tag)
else:
storage = desc.storage
tag_name = _define_local_scalar(pv, sdfg, state, dace.int32, storage)
tag_node = state.add_access(tag_name)
tag_tasklet = state.add_tasklet('_set_tag_', {}, {'__out'},
'__out = {}'.format(tag))
state.add_edge(tag_tasklet, '__out', tag_node, None,
Memlet.simple(tag_name, '0'))
if buf_range:
buf_mem = Memlet.simple(buf_name, buf_range)
else:
buf_mem = Memlet.from_array(buf_name, desc)
if req_range:
req_mem = Memlet.simple(req_name, req_range)
else:
req_mem = Memlet.from_array(req_name, req_desc)
if dst_range:
dst_mem = Memlet.simple(dst_name, dst_range)
else:
dst_mem = Memlet.simple(dst_name, '0')
if tag_range:
tag_mem = Memlet.simple(tag_name, tag_range)
else:
tag_mem = Memlet.simple(tag_name, '0')
state.add_edge(buf_node, None, libnode, '_buffer', buf_mem)
state.add_edge(dst_node, None, libnode, '_dest', dst_mem)
state.add_edge(tag_node, None, libnode, '_tag', tag_mem)
state.add_edge(libnode, '_request', req_node, None, req_mem)
return None
def _irecv(pv: 'ProgramVisitor', sdfg: SDFG, state: SDFGState, buffer: str,
src: Union[str, sp.Expr, Number], tag: Union[str, sp.Expr,
Number], request: str):
from dace.libraries.mpi.nodes.irecv import Irecv
libnode = Irecv('_Irecv_')
buf_range = None
if isinstance(buffer, tuple):
buf_name, buf_range = buffer
else:
buf_name = buffer
desc = sdfg.arrays[buf_name]
buf_node = state.add_read(buf_name)
req_range = None
if isinstance(request, tuple):
req_name, req_range = request
else:
req_name = request
req_desc = sdfg.arrays[req_name]
req_node = state.add_write(req_name)
conn = libnode.out_connectors
conn = {
c: (dtypes.pointer(desc.dtype) if c == '_buffer' else t)
for c, t in conn.items()
}
conn = {
c: (dtypes.pointer(req_desc.dtype) if c == '_request' else t)
for c, t in conn.items()
}
libnode.out_connectors = conn
src_range = None
if isinstance(src, tuple):
src_name, src_range = src
src_node = state.add_read(src_name)
elif isinstance(src, str) and src in sdfg.arrays.keys():
src_name = src
src_node = state.add_read(src_name)
else:
storage = desc.storage
src_name = _define_local_scalar(pv, sdfg, state, dace.int32, storage)
src_node = state.add_access(src_name)
src_tasklet = state.add_tasklet('_set_src_', {}, {'__out'},
'__out = {}'.format(src))
state.add_edge(src_tasklet, '__out', src_node, None,
Memlet.simple(src_name, '0'))
tag_range = None
if isinstance(tag, tuple):
tag_name, tag_range = tag
tag_node = state.add_read(tag_name)
if isinstance(tag, str) and tag in sdfg.arrays.keys():
tag_name = tag
tag_node = state.add_read(tag)
else:
storage = desc.storage
tag_name = _define_local_scalar(pv, sdfg, state, dace.int32, storage)
tag_node = state.add_access(tag_name)
tag_tasklet = state.add_tasklet('_set_tag_', {}, {'__out'},
'__out = {}'.format(tag))
state.add_edge(tag_tasklet, '__out', tag_node, None,
Memlet.simple(tag_name, '0'))
if buf_range:
buf_mem = Memlet.simple(buf_name, buf_range)
else:
buf_mem = Memlet.from_array(buf_name, desc)
if req_range:
req_mem = Memlet.simple(req_name, req_range)
else:
req_mem = Memlet.from_array(req_name, req_desc)
if src_range:
src_mem = Memlet.simple(src_name, src_range)
else:
src_mem = Memlet.simple(src_name, '0')
if tag_range:
tag_mem = Memlet.simple(tag_name, tag_range)
else:
tag_mem = Memlet.simple(tag_name, '0')
state.add_edge(libnode, '_buffer', buf_node, None, buf_mem)
state.add_edge(src_node, None, libnode, '_src', src_mem)
state.add_edge(tag_node, None, libnode, '_tag', tag_mem)
state.add_edge(libnode, '_request', req_node, None, req_mem)
return None
def can_be_applied(self,
state: SDFGState,
candidate,
expr_index,
sdfg,
strict=False):
nested_sdfg = self.nested_sdfg(sdfg)
if nested_sdfg.no_inline:
return False
# Ensure the state only contains a nested SDFG and input/output access
# nodes
for node in state.nodes():
if isinstance(node, nodes.NestedSDFG):
if node is not nested_sdfg:
return False
elif isinstance(node, nodes.AccessNode):
# Must be connected to nested SDFG
# if nested_sdfg in state.predecessors(nested_sdfg):
# if state.in_degree(node) > 0:
# return False
found = False
for e in state.out_edges(node):
if e.dst is not nested_sdfg:
return False
if state.in_degree(node) > 0:
return False
# Only accept full ranges for now. TODO(later): Improve
if e.data.subset != subsets.Range.from_array(
sdfg.arrays[node.data]):
return False
# Do not accept views. TODO(later): Improve
outer_desc = sdfg.arrays[node.data]
inner_desc = nested_sdfg.sdfg.arrays[e.dst_conn]
if (outer_desc.shape != inner_desc.shape
or outer_desc.strides != inner_desc.strides):
return False
found = True
for e in state.in_edges(node):
if e.src is not nested_sdfg:
return False
if state.out_degree(node) > 0:
return False
# Only accept full ranges for now. TODO(later): Improve
if e.data.subset != subsets.Range.from_array(
sdfg.arrays[node.data]):
return False
# Do not accept views. TODO(later): Improve
outer_desc = sdfg.arrays[node.data]
inner_desc = nested_sdfg.sdfg.arrays[e.src_conn]
if (outer_desc.shape != inner_desc.shape
or outer_desc.strides != inner_desc.strides):
return False
found = True
# elif nested_sdfg in state.successors(nested_sdfg):
# if state.out_degree(node) > 0:
# return False
if not found:
return False
else:
return False
return True
def apply(self, state: SDFGState, sdfg: SDFG) -> nodes.AccessNode:
access: nodes.AccessNode = self.access
# Get memlet paths
first_edge = state.in_edges(access)[0]
second_edge = state.out_edges(access)[0]
first_mpath = state.memlet_path(first_edge)
second_mpath = state.memlet_path(second_edge)
# Create new stream of shape 1
desc = sdfg.arrays[access.data]
# Qualify name to avoid name clashes if memory interfaces are not decoupled for Xilinx
stream_name = "stream_" + access.data
name, newdesc = sdfg.add_stream(stream_name,
desc.dtype,
buffer_size=self.buffer_size,
storage=self.storage,
transient=True,
find_new_name=True)
# Remove transient array if possible
for ostate in sdfg.nodes():
if ostate is state:
continue
if any(n.data == access.data for n in ostate.data_nodes()):
break
else:
if desc.transient:
del sdfg.arrays[access.data]
# Replace memlets in path with stream access
for e in first_mpath:
e.data = mm.Memlet(data=name, subset='0')
if isinstance(e.src, nodes.NestedSDFG):
e.data.dynamic = True
_streamify_recursive(e.src, e.src_conn, newdesc)
if isinstance(e.dst, nodes.NestedSDFG):
e.data.dynamic = True
_streamify_recursive(e.dst, e.dst_conn, newdesc)
for e in second_mpath:
e.data = mm.Memlet(data=name, subset='0')
if isinstance(e.src, nodes.NestedSDFG):
e.data.dynamic = True
_streamify_recursive(e.src, e.src_conn, newdesc)
if isinstance(e.dst, nodes.NestedSDFG):
e.data.dynamic = True
_streamify_recursive(e.dst, e.dst_conn, newdesc)
# Replace array access node with two stream access nodes
wnode = state.add_write(name)
rnode = state.add_read(name)
state.remove_edge(first_edge)
state.add_edge(first_edge.src, first_edge.src_conn, wnode,
first_edge.dst_conn, first_edge.data)
state.remove_edge(second_edge)
state.add_edge(rnode, second_edge.src_conn, second_edge.dst,
second_edge.dst_conn, second_edge.data)
# Remove original access node
state.remove_node(access)
return wnode, rnode
def apply(self, state: SDFGState, sdfg: SDFG) -> nodes.AccessNode:
dnode: nodes.AccessNode = self.access
if self.expr_index == 0:
edges = state.out_edges(dnode)
else:
edges = state.in_edges(dnode)
# To understand how many components we need to create, all map ranges
# throughout memlet paths must match exactly. We thus create a
# dictionary of unique ranges
mapping: Dict[Tuple[subsets.Range],
List[gr.MultiConnectorEdge[mm.Memlet]]] = defaultdict(
list)
ranges = {}
for edge in edges:
mpath = state.memlet_path(edge)
ranges[edge] = _collect_map_ranges(state, mpath)
mapping[tuple(r[1] for r in ranges[edge])].append(edge)
# Collect all edges with the same memory access pattern
components_to_create: Dict[
Tuple[symbolic.SymbolicType],
List[gr.MultiConnectorEdge[mm.Memlet]]] = defaultdict(list)
for edges_with_same_range in mapping.values():
for edge in edges_with_same_range:
# Get memlet path and innermost edge
mpath = state.memlet_path(edge)
innermost_edge = copy.deepcopy(mpath[-1] if self.expr_index ==
0 else mpath[0])
# Store memlets of the same access in the same component
expr = _canonicalize_memlet(innermost_edge.data, ranges[edge])
components_to_create[expr].append((innermost_edge, edge))
components = list(components_to_create.values())
# Split out components that have dependencies between them to avoid
# deadlocks
if self.expr_index == 0:
ccs_to_add = []
for i, component in enumerate(components):
edges_to_remove = set()
for cedge in component:
if any(
nx.has_path(state.nx, o[1].dst, cedge[1].dst)
for o in component if o is not cedge):
ccs_to_add.append([cedge])
edges_to_remove.add(cedge)
if edges_to_remove:
components[i] = [
c for c in component if c not in edges_to_remove
]
components.extend(ccs_to_add)
# End of split
desc = sdfg.arrays[dnode.data]
# Create new streams of shape 1
streams = {}
mpaths = {}
for edge in edges:
if self.use_memory_buffering:
arrname = str(self.access)
# Add gearbox
total_size = edge.data.volume
vector_size = int(self.memory_buffering_target_bytes /
desc.dtype.bytes)
if not is_int(sdfg.arrays[dnode.data].shape[-1]):
warnings.warn(
"Using the MemoryBuffering transformation is potential unsafe since {sym} is not an integer. There should be no issue if {sym} % {vec} == 0"
.format(sym=sdfg.arrays[dnode.data].shape[-1],
vec=vector_size))
for i in sdfg.arrays[dnode.data].strides:
if not is_int(i):
warnings.warn(
"Using the MemoryBuffering transformation is potential unsafe since {sym} is not an integer. There should be no issue if {sym} % {vec} == 0"
.format(sym=i, vec=vector_size))
if self.expr_index == 0: # Read
edges = state.out_edges(dnode)
gearbox_input_type = dtypes.vector(desc.dtype, vector_size)
gearbox_output_type = desc.dtype
gearbox_read_volume = total_size / vector_size
gearbox_write_volume = total_size
else: # Write
edges = state.in_edges(dnode)
gearbox_input_type = desc.dtype
gearbox_output_type = dtypes.vector(
desc.dtype, vector_size)
gearbox_read_volume = total_size
gearbox_write_volume = total_size / vector_size
input_gearbox_name, input_gearbox_newdesc = sdfg.add_stream(
"gearbox_input",
gearbox_input_type,
buffer_size=self.buffer_size,
storage=self.storage,
transient=True,
find_new_name=True)
output_gearbox_name, output_gearbox_newdesc = sdfg.add_stream(
"gearbox_output",
gearbox_output_type,
buffer_size=self.buffer_size,
storage=self.storage,
transient=True,
find_new_name=True)
read_to_gearbox = state.add_read(input_gearbox_name)
write_from_gearbox = state.add_write(output_gearbox_name)
gearbox = Gearbox(total_size / vector_size)
state.add_node(gearbox)
state.add_memlet_path(read_to_gearbox,
gearbox,
dst_conn="from_memory",
memlet=Memlet(
input_gearbox_name + "[0]",
volume=gearbox_read_volume))
state.add_memlet_path(gearbox,
write_from_gearbox,
src_conn="to_kernel",
memlet=Memlet(
output_gearbox_name + "[0]",
volume=gearbox_write_volume))
if self.expr_index == 0:
streams[edge] = input_gearbox_name
name = output_gearbox_name
newdesc = output_gearbox_newdesc
else:
streams[edge] = output_gearbox_name
name = input_gearbox_name
newdesc = input_gearbox_newdesc
else:
# Qualify name to avoid name clashes if memory interfaces are not decoupled for Xilinx
stream_name = "stream_" + dnode.data
name, newdesc = sdfg.add_stream(stream_name,
desc.dtype,
buffer_size=self.buffer_size,
storage=self.storage,
transient=True,
find_new_name=True)
streams[edge] = name
# Add these such that we can easily use output_gearbox_name and input_gearbox_name without using if statements
output_gearbox_name = name
input_gearbox_name = name
mpath = state.memlet_path(edge)
mpaths[edge] = mpath
# Replace memlets in path with stream access
for e in mpath:
e.data = mm.Memlet(data=name,
subset='0',
other_subset=e.data.other_subset)
if isinstance(e.src, nodes.NestedSDFG):
e.data.dynamic = True
_streamify_recursive(e.src, e.src_conn, newdesc)
if isinstance(e.dst, nodes.NestedSDFG):
e.data.dynamic = True
_streamify_recursive(e.dst, e.dst_conn, newdesc)
# Replace access node and memlet tree with one access
if self.expr_index == 0:
replacement = state.add_read(output_gearbox_name)
state.remove_edge(edge)
state.add_edge(replacement, edge.src_conn, edge.dst,
edge.dst_conn, edge.data)
else:
replacement = state.add_write(input_gearbox_name)
state.remove_edge(edge)
state.add_edge(edge.src, edge.src_conn, replacement,
edge.dst_conn, edge.data)
if self.use_memory_buffering:
arrname = str(self.access)
vector_size = int(self.memory_buffering_target_bytes /
desc.dtype.bytes)
# Vectorize access to global array.
dtype = sdfg.arrays[arrname].dtype
sdfg.arrays[arrname].dtype = dtypes.vector(dtype, vector_size)
new_shape = list(sdfg.arrays[arrname].shape)
contigidx = sdfg.arrays[arrname].strides.index(1)
new_shape[contigidx] /= vector_size
try:
new_shape[contigidx] = int(new_shape[contigidx])
except TypeError:
pass
sdfg.arrays[arrname].shape = new_shape
# Change strides
new_strides: List = list(sdfg.arrays[arrname].strides)
for i in range(len(new_strides)):
if i == len(new_strides
) - 1: # Skip last dimension since it is always 1
continue
new_strides[i] = new_strides[i] / vector_size
sdfg.arrays[arrname].strides = new_strides
post_state = get_post_state(sdfg, state)
if post_state != None:
# Change subset in the post state such that the correct amount of memory is copied back from the device
for e in post_state.edges():
if e.data.data == self.access.data:
new_subset = list(e.data.subset)
i, j, k = new_subset[-1]
new_subset[-1] = (i, (j + 1) / vector_size - 1, k)
e.data = mm.Memlet(data=str(e.src),
subset=subsets.Range(new_subset))
# Make read/write components
ionodes = []
for component in components:
# Pick the first edge as the edge to make the component from
innermost_edge, outermost_edge = component[0]
mpath = mpaths[outermost_edge]
mapname = streams[outermost_edge]
innermost_edge.data.other_subset = None
# Get edge data and streams
if self.expr_index == 0:
opname = 'read'
path = [e.dst for e in mpath[:-1]]
rmemlets = [(dnode, '__inp', innermost_edge.data)]
wmemlets = []
for i, (_, edge) in enumerate(component):
name = streams[edge]
ionode = state.add_write(name)
ionodes.append(ionode)
wmemlets.append(
(ionode, '__out%d' % i, mm.Memlet(data=name,
subset='0')))
code = '\n'.join('__out%d = __inp' % i
for i in range(len(component)))
else:
# More than one input stream might mean a data race, so we only
# address the first one in the tasklet code
if len(component) > 1:
warnings.warn(
f'More than one input found for the same index for {dnode.data}'
)
opname = 'write'
path = [state.entry_node(e.src) for e in reversed(mpath[1:])]
wmemlets = [(dnode, '__out', innermost_edge.data)]
rmemlets = []
for i, (_, edge) in enumerate(component):
name = streams[edge]
ionode = state.add_read(name)
ionodes.append(ionode)
rmemlets.append(
(ionode, '__inp%d' % i, mm.Memlet(data=name,
subset='0')))
code = '__out = __inp0'
# Create map structure for read/write component
maps = []
for entry in path:
map: nodes.Map = entry.map
ranges = [(p, (r[0], r[1], r[2]))
for p, r in zip(map.params, map.range)]
# Change ranges of map
if self.use_memory_buffering:
# Find edges from/to map
edge_subset = [
a_tuple[0]
for a_tuple in list(innermost_edge.data.subset)
]
# Change range of map
if isinstance(edge_subset[-1], symbol) and str(
edge_subset[-1]) == map.params[-1]:
if not is_int(ranges[-1][1][1]):
warnings.warn(
"Using the MemoryBuffering transformation is potential unsafe since {sym} is not an integer. There should be no issue if {sym} % {vec} == 0"
.format(sym=ranges[-1][1][1].args[1],
vec=vector_size))
ranges[-1] = (ranges[-1][0],
(ranges[-1][1][0],
(ranges[-1][1][1] + 1) / vector_size -
1, ranges[-1][1][2]))
elif isinstance(edge_subset[-1], sympy.core.add.Add):
for arg in edge_subset[-1].args:
if isinstance(
arg,
symbol) and str(arg) == map.params[-1]:
if not is_int(ranges[-1][1][1]):
warnings.warn(
"Using the MemoryBuffering transformation is potential unsafe since {sym} is not an integer. There should be no issue if {sym} % {vec} == 0"
.format(sym=ranges[-1][1][1].args[1],
vec=vector_size))
ranges[-1] = (ranges[-1][0], (
ranges[-1][1][0],
(ranges[-1][1][1] + 1) / vector_size - 1,
ranges[-1][1][2]))
maps.append(
state.add_map(f'__s{opname}_{mapname}', ranges,
map.schedule))
tasklet = state.add_tasklet(
f'{opname}_{mapname}',
{m[1]
for m in rmemlets},
{m[1]
for m in wmemlets},
code,
)
for node, cname, memlet in rmemlets:
state.add_memlet_path(node,
*(me for me, _ in maps),
tasklet,
dst_conn=cname,
memlet=memlet)
for node, cname, memlet in wmemlets:
state.add_memlet_path(tasklet,
*(mx for _, mx in reversed(maps)),
node,
src_conn=cname,
memlet=memlet)
return ionodes
def _elementwise(sdfg: SDFG,
state: SDFGState,
func: str,
in_array: str,
out_array=None):
"""Apply a lambda function to each element in the input"""
inparr = sdfg.arrays[in_array]
restype = sdfg.arrays[in_array].dtype
if out_array is None:
out_array, outarr = sdfg.add_temp_transient(inparr.shape, restype,
inparr.storage)
else:
outarr = sdfg.arrays[out_array]
func_ast = ast.parse(func)
try:
lambda_ast = func_ast.body[0].value
if len(lambda_ast.args.args) != 1:
raise SyntaxError(
"Expected lambda with one arg, but {} has {}".format(
func, len(lambda_ast.args.arrgs)))
arg = lambda_ast.args.args[0].arg
body = astutils.unparse(lambda_ast.body)
except AttributeError:
raise SyntaxError("Could not parse func {}".format(func))
code = "__out = {}".format(body)
num_elements = reduce(lambda x, y: x * y, inparr.shape)
if num_elements == 1:
inp = state.add_read(in_array)
out = state.add_write(out_array)
tasklet = state.add_tasklet("_elementwise_", {arg}, {'__out'}, code)
state.add_edge(inp, None, tasklet, arg,
Memlet.from_array(in_array, inparr))
state.add_edge(tasklet, '__out', out, None,
Memlet.from_array(out_array, outarr))
else:
state.add_mapped_tasklet(
name="_elementwise_",
map_ranges={
'__i%d' % i: '0:%s' % n
for i, n in enumerate(inparr.shape)
},
inputs={
arg:
Memlet.simple(
in_array,
','.join(['__i%d' % i for i in range(len(inparr.shape))]))
},
code=code,
outputs={
'__out':
Memlet.simple(
out_array,
','.join(['__i%d' % i for i in range(len(inparr.shape))]))
},
external_edges=True)
return out_array
def _gather(pv: 'ProgramVisitor',
sdfg: SDFG,
state: SDFGState,
in_buffer: str,
out_buffer: str,
root: Union[str, sp.Expr, Number] = 0):
from dace.libraries.mpi.nodes.gather import Gather
libnode = Gather('_Gather_')
in_desc = sdfg.arrays[in_buffer]
out_desc = sdfg.arrays[out_buffer]
in_node = state.add_read(in_buffer)
out_node = state.add_write(out_buffer)
if isinstance(root, str) and root in sdfg.arrays.keys():
root_node = state.add_read(root)
else:
storage = in_desc.storage
root_name = _define_local_scalar(pv, sdfg, state, dace.int32, storage)
root_node = state.add_access(root_name)
root_tasklet = state.add_tasklet('_set_root_', {}, {'__out'},
'__out = {}'.format(root))
state.add_edge(root_tasklet, '__out', root_node, None,
Memlet.simple(root_name, '0'))
state.add_edge(in_node, None, libnode, '_inbuffer',
Memlet.from_array(in_buffer, in_desc))
state.add_edge(root_node, None, libnode, '_root',
Memlet.simple(root_node.data, '0'))
state.add_edge(libnode, '_outbuffer', out_node, None,
Memlet.from_array(out_buffer, out_desc))
return None
def _matmult(visitor, sdfg: SDFG, state: SDFGState, op1: str, op2: str):
from dace.libraries.blas.nodes.matmul import MatMul # Avoid import loop
arr1 = sdfg.arrays[op1]
arr2 = sdfg.arrays[op2]
if len(arr1.shape) > 1 and len(arr2.shape) > 1: # matrix * matrix
if len(arr1.shape) > 3 or len(arr2.shape) > 3:
raise SyntaxError(
'Matrix multiplication of tensors of dimensions > 3 '
'not supported')
if arr1.shape[-1] != arr2.shape[-2]:
raise SyntaxError('Matrix dimension mismatch %s != %s' %
(arr1.shape[-1], arr2.shape[-2]))
from dace.libraries.blas.nodes.matmul import _get_batchmm_opts
# Determine batched multiplication
bopt = _get_batchmm_opts(arr1.shape, arr1.strides, arr2.shape,
arr2.strides, None, None)
if bopt:
output_shape = (bopt['b'], arr1.shape[-2], arr2.shape[-1])
else:
output_shape = (arr1.shape[-2], arr2.shape[-1])
elif len(arr1.shape) == 2 and len(arr2.shape) == 1: # matrix * vector
if arr1.shape[1] != arr2.shape[0]:
raise SyntaxError("Number of matrix columns {} must match"
"size of vector {}.".format(
arr1.shape[1], arr2.shape[0]))
output_shape = (arr1.shape[0], )
elif len(arr1.shape) == 1 and len(arr2.shape) == 1: # vector * vector
if arr1.shape[0] != arr2.shape[0]:
raise SyntaxError("Vectors in vector product must have same size: "
"{} vs. {}".format(arr1.shape[0], arr2.shape[0]))
output_shape = (1, )
else: # Dunno what this is, bail
raise SyntaxError(
"Cannot multiply arrays with shapes: {} and {}".format(
arr1.shape, arr2.shape))
type1 = arr1.dtype.type
type2 = arr2.dtype.type
restype = dace.DTYPE_TO_TYPECLASS[np.result_type(type1, type2).type]
op3, arr3 = sdfg.add_temp_transient(output_shape, restype, arr1.storage)
acc1 = state.add_read(op1)
acc2 = state.add_read(op2)
acc3 = state.add_write(op3)
tasklet = MatMul('_MatMult_', restype)
state.add_node(tasklet)
state.add_edge(acc1, None, tasklet, '_a', dace.Memlet.from_array(op1, arr1))
state.add_edge(acc2, None, tasklet, '_b', dace.Memlet.from_array(op2, arr2))
state.add_edge(tasklet, '_c', acc3, None, dace.Memlet.from_array(op3, arr3))
return op3
def nccl_allreduce(pv: 'ProgramVisitor',
sdfg: SDFG,
state: SDFGState,
redfunction: Callable[[Any, Any], Any],
in_buffer: str,
out_buffer: Union[str, None] = None,
group_handle: str = None):
inputs = {"_inbuffer"}
outputs = {"_outbuffer"}
if isinstance(group_handle, str):
gh_start = False
if group_handle in sdfg.arrays.keys():
gh_name = group_handle
gh_out = state.add_access(gh_name)
gh_in = state.add_access(gh_name)
inputs.add("_group_handle")
else:
gh_start = True
gh_name = _define_local_scalar(pv, sdfg, state, dace.int32,
dtypes.StorageType.GPU_Global)
gh_out = state.add_access(gh_name)
outputs.add("_group_handle")
libnode = Allreduce(inputs=inputs, outputs=outputs, wcr=redfunction)
if isinstance(group_handle, str):
gh_memlet = Memlet.simple(gh_name, '0')
if not gh_start:
state.add_edge(gh_in, None, libnode, "_group_handle", gh_memlet)
state.add_edge(libnode, "_group_handle", gh_out, None, gh_memlet)
# If out_buffer is not specified, the operation will be in-place.
if out_buffer is None:
out_buffer = in_buffer
# Add nodes
in_node = state.add_read(in_buffer)
out_node = state.add_write(out_buffer)
# Connect nodes
state.add_edge(in_node, None, libnode, '_inbuffer', Memlet(in_buffer))
state.add_edge(libnode, '_outbuffer', out_node, None, Memlet(out_buffer))
return []
def expansion(node: 'Reduce', state: SDFGState, sdfg: SDFG):
node.validate(sdfg, state)
inedge: graph.MultiConnectorEdge = state.in_edges(node)[0]
outedge: graph.MultiConnectorEdge = state.out_edges(node)[0]
input_dims = len(inedge.data.subset)
output_dims = len(outedge.data.subset)
input_data = sdfg.arrays[inedge.data.data]
output_data = sdfg.arrays[outedge.data.data]
# Get reduction type for OpenMP
redtype = detect_reduction_type(node.wcr, openmp=True)
if redtype not in ExpandReduceOpenMP._REDUCTION_TYPE_TO_OPENMP:
warnings.warn('Reduction type not supported for "%s"' % node.wcr)
return ExpandReducePure.expansion(node, state, sdfg)
omptype, expr = ExpandReduceOpenMP._REDUCTION_TYPE_TO_OPENMP[redtype]
# Standardize axes
axes = node.axes if node.axes else [i for i in range(input_dims)]
outer_loops = len(axes) != input_dims
# Create OpenMP clause
if outer_loops:
code = '#pragma omp parallel for collapse({cdim})\n'.format(
cdim=output_dims)
else:
code = ''
from dace.codegen.targets.cpp import sym2cpp
# Output loops
out_offset = []
if outer_loops:
for i, sz in enumerate(outedge.data.subset.size()):
code += 'for (int _o{i} = 0; _o{i} < {sz}; ++_o{i}) {{\n'.format(
i=i, sz=sym2cpp(sz))
out_offset.append('_o%d * %s' %
(i, sym2cpp(output_data.strides[i])))
else:
out_offset.append('0')
outexpr = '_out[%s]' % ' + '.join(out_offset)
# Write identity value first
if node.identity is not None:
code += '%s = %s;\n' % (outexpr, sym2cpp(node.identity))
# Reduction OpenMP clause
code += '#pragma omp parallel for collapse({cdim}) ' \
'reduction({rtype}: {oexpr})\n'.format(cdim=len(axes), rtype=omptype,
oexpr=outexpr)
# Reduction loops
for i, axis in enumerate(sorted(axes)):
sz = sym2cpp(inedge.data.subset.size()[axis])
code += 'for (int _i{i} = 0; _i{i} < {sz}; ++_i{i}) {{\n'.format(
i=i, sz=sz)
# Prepare input offset expression
in_offset = []
ictr, octr = 0, 0
for i in range(input_dims):
if i in axes:
result = '_i%d' % ictr
ictr += 1
else:
result = '_o%d' % octr
octr += 1
in_offset.append('%s * %s' %
(result, sym2cpp(input_data.strides[i])))
in_offset = ' + '.join(in_offset)
# Reduction expression
code += expr.format(i='_in[%s]' % in_offset, o=outexpr)
code += '\n'
# Closing braces
code += '}\n' * len(axes)
if outer_loops:
code += '}\n' * output_dims
# Make tasklet
tnode = dace.nodes.Tasklet('reduce',
{'_in': dace.pointer(input_data.dtype)},
{'_out': dace.pointer(output_data.dtype)},
code,
language=dace.Language.CPP)
# Rename outer connectors and add to node
inedge._dst_conn = '_in'
outedge._src_conn = '_out'
node.add_in_connector('_in')
node.add_out_connector('_out')
return tnode
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)
# 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 = {}
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, 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 expansion(node: 'Reduce', state: SDFGState, sdfg: SDFG):
from dace.codegen.prettycode import CodeIOStream
from dace.codegen.targets.cpp import unparse_cr_split, cpp_array_expr
node.validate(sdfg, state)
input_edge: graph.MultiConnectorEdge = state.in_edges(node)[0]
output_edge: graph.MultiConnectorEdge = state.out_edges(node)[0]
input_dims = len(input_edge.data.subset)
output_dims = len(output_edge.data.subset)
input_data = sdfg.arrays[input_edge.data.data]
output_data = sdfg.arrays[output_edge.data.data]
# Setup all locations in which code will be written
cuda_globalcode = CodeIOStream()
cuda_initcode = CodeIOStream()
cuda_exitcode = CodeIOStream()
host_globalcode = CodeIOStream()
host_localcode = CodeIOStream()
output_memlet = output_edge.data
# Try to autodetect reduction type
redtype = detect_reduction_type(node.wcr)
node_id = state.node_id(node)
state_id = sdfg.node_id(state)
idstr = '{sdfg}_{state}_{node}'.format(sdfg=sdfg.name,
state=state_id,
node=node_id)
if node.out_connectors:
dtype = next(node.out_connectors.values())
else:
dtype = sdfg.arrays[output_memlet.data].dtype
output_type = dtype.ctype
if node.identity is None:
raise ValueError('For device reduce nodes, initial value must be '
'specified')
# Create a functor or use an existing one for reduction
if redtype == dtypes.ReductionType.Custom:
body, [arg1, arg2] = unparse_cr_split(sdfg, node.wcr)
cuda_globalcode.write(
"""
struct __reduce_{id} {{
template <typename T>
DACE_HDFI T operator()(const T &{arg1}, const T &{arg2}) const {{
{contents}
}}
}};""".format(id=idstr, arg1=arg1, arg2=arg2, contents=body), sdfg,
state_id, node_id)
reduce_op = ', __reduce_' + idstr + '(), ' + symstr(node.identity)
elif redtype in ExpandReduceCUDADevice._SPECIAL_RTYPES:
reduce_op = ''
else:
credtype = 'dace::ReductionType::' + str(
redtype)[str(redtype).find('.') + 1:]
reduce_op = ((', dace::_wcr_fixed<%s, %s>()' %
(credtype, output_type)) + ', ' +
symstr(node.identity))
# Obtain some SDFG-related information
input_memlet = input_edge.data
reduce_shape = input_memlet.subset.bounding_box_size()
num_items = ' * '.join(symstr(s) for s in reduce_shape)
overapprox_memlet = dcpy(input_memlet)
if any(
str(s) not in sdfg.free_symbols.union(sdfg.constants.keys())
for s in overapprox_memlet.subset.free_symbols):
propagation.propagate_states(sdfg)
for p, r in state.ranges.items():
overapprox_memlet = propagation.propagate_subset(
[overapprox_memlet], input_data, [p], r)
overapprox_shape = overapprox_memlet.subset.bounding_box_size()
overapprox_items = ' * '.join(symstr(s) for s in overapprox_shape)
input_dims = input_memlet.subset.dims()
output_dims = output_memlet.subset.data_dims()
reduce_all_axes = (node.axes is None or len(node.axes) == input_dims)
if reduce_all_axes:
reduce_last_axes = False
else:
reduce_last_axes = sorted(node.axes) == list(
range(input_dims - len(node.axes), input_dims))
if not reduce_all_axes and not reduce_last_axes:
warnings.warn(
'Multiple axis reductions not supported with this expansion. '
'Falling back to the pure expansion.')
return ExpandReducePureSequentialDim.expansion(node, state, sdfg)
# Verify that data is on the GPU
if input_data.storage not in [
dtypes.StorageType.GPU_Global, dtypes.StorageType.CPU_Pinned
]:
warnings.warn('Input of GPU reduction must either reside '
' in global GPU memory or pinned CPU memory')
return ExpandReducePure.expansion(node, state, sdfg)
if output_data.storage not in [
dtypes.StorageType.GPU_Global, dtypes.StorageType.CPU_Pinned
]:
warnings.warn('Output of GPU reduction must either reside '
' in global GPU memory or pinned CPU memory')
return ExpandReducePure.expansion(node, state, sdfg)
# Determine reduction type
kname = (ExpandReduceCUDADevice._SPECIAL_RTYPES[redtype] if redtype
in ExpandReduceCUDADevice._SPECIAL_RTYPES else 'Reduce')
# Create temp memory for this GPU
cuda_globalcode.write(
"""
void *__cub_storage_{sdfg}_{state}_{node} = NULL;
size_t __cub_ssize_{sdfg}_{state}_{node} = 0;
""".format(sdfg=sdfg.name, state=state_id, node=node_id), sdfg,
state_id, node)
if reduce_all_axes:
reduce_type = 'DeviceReduce'
reduce_range = overapprox_items
reduce_range_def = 'size_t num_items'
reduce_range_use = 'num_items'
reduce_range_call = num_items
elif reduce_last_axes:
num_reduce_axes = len(node.axes)
not_reduce_axes = reduce_shape[:-num_reduce_axes]
reduce_axes = reduce_shape[-num_reduce_axes:]
overapprox_not_reduce_axes = overapprox_shape[:-num_reduce_axes]
overapprox_reduce_axes = overapprox_shape[-num_reduce_axes:]
num_segments = ' * '.join([symstr(s) for s in not_reduce_axes])
segment_size = ' * '.join([symstr(s) for s in reduce_axes])
overapprox_num_segments = ' * '.join(
[symstr(s) for s in overapprox_not_reduce_axes])
overapprox_segment_size = ' * '.join(
[symstr(s) for s in overapprox_reduce_axes])
reduce_type = 'DeviceSegmentedReduce'
iterator = 'dace::stridedIterator({size})'.format(
size=overapprox_segment_size)
reduce_range = '{num}, {it}, {it} + 1'.format(
num=overapprox_num_segments, it=iterator)
reduce_range_def = 'size_t num_segments, size_t segment_size'
iterator_use = 'dace::stridedIterator(segment_size)'
reduce_range_use = 'num_segments, {it}, {it} + 1'.format(
it=iterator_use)
reduce_range_call = '%s, %s' % (num_segments, segment_size)
# Call CUB to get the storage size, allocate and free it
cuda_initcode.write(
"""
cub::{reduce_type}::{kname}(nullptr, __cub_ssize_{sdfg}_{state}_{node},
({intype}*)nullptr, ({outtype}*)nullptr, {reduce_range}{redop});
cudaMalloc(&__cub_storage_{sdfg}_{state}_{node}, __cub_ssize_{sdfg}_{state}_{node});
""".format(sdfg=sdfg.name,
state=state_id,
node=node_id,
reduce_type=reduce_type,
reduce_range=reduce_range,
redop=reduce_op,
intype=input_data.dtype.ctype,
outtype=output_data.dtype.ctype,
kname=kname), sdfg, state_id, node)
cuda_exitcode.write(
'cudaFree(__cub_storage_{sdfg}_{state}_{node});'.format(
sdfg=sdfg.name, state=state_id, node=node_id), sdfg, state_id,
node)
# Write reduction function definition
cuda_globalcode.write("""
DACE_EXPORTED void __dace_reduce_{id}({intype} *input, {outtype} *output, {reduce_range_def}, cudaStream_t stream);
void __dace_reduce_{id}({intype} *input, {outtype} *output, {reduce_range_def}, cudaStream_t stream)
{{
cub::{reduce_type}::{kname}(__cub_storage_{id}, __cub_ssize_{id},
input, output, {reduce_range_use}{redop}, stream);
}}
""".format(id=idstr,
intype=input_data.dtype.ctype,
outtype=output_data.dtype.ctype,
reduce_type=reduce_type,
reduce_range_def=reduce_range_def,
reduce_range_use=reduce_range_use,
kname=kname,
redop=reduce_op))
# Write reduction function definition in caller file
host_globalcode.write(
"""
DACE_EXPORTED void __dace_reduce_{id}({intype} *input, {outtype} *output, {reduce_range_def}, cudaStream_t stream);
""".format(id=idstr,
reduce_range_def=reduce_range_def,
intype=input_data.dtype.ctype,
outtype=output_data.dtype.ctype), sdfg, state_id, node)
# Call reduction function where necessary
host_localcode.write(
'__dace_reduce_{id}(_in, _out, {reduce_range_call}, __dace_current_stream);'
.format(id=idstr, reduce_range_call=reduce_range_call))
# Make tasklet
tnode = dace.nodes.Tasklet('reduce',
{'_in': dace.pointer(input_data.dtype)},
{'_out': dace.pointer(output_data.dtype)},
host_localcode.getvalue(),
language=dace.Language.CPP)
# Add the rest of the code
sdfg.append_global_code(host_globalcode.getvalue())
sdfg.append_global_code(cuda_globalcode.getvalue(), 'cuda')
sdfg.append_init_code(cuda_initcode.getvalue(), 'cuda')
sdfg.append_exit_code(cuda_exitcode.getvalue(), 'cuda')
# Rename outer connectors and add to node
input_edge._dst_conn = '_in'
output_edge._src_conn = '_out'
node.add_in_connector('_in')
node.add_out_connector('_out')
return tnode
def replicate_scope(sdfg: SDFG, state: SDFGState,
scope: ScopeSubgraphView) -> ScopeSubgraphView:
"""
Replicates a scope subgraph view within a state, reconnecting all external
edges to the same nodes.
:param sdfg: The SDFG in which the subgraph scope resides.
:param state: The SDFG state in which the subgraph scope resides.
:param scope: The scope subgraph to replicate.
:return: A reconnected replica of the scope.
"""
exit_node = state.exit_node(scope.entry)
# Replicate internal graph
new_nodes = []
new_entry = None
new_exit = None
for node in scope.nodes():
node_copy = copy.deepcopy(node)
if node == scope.entry:
new_entry = node_copy
elif node == exit_node:
new_exit = node_copy
state.add_node(node_copy)
new_nodes.append(node_copy)
for edge in scope.edges():
src = scope.nodes().index(edge.src)
dst = scope.nodes().index(edge.dst)
state.add_edge(new_nodes[src], edge.src_conn, new_nodes[dst],
edge.dst_conn, copy.deepcopy(edge.data))
# Reconnect external scope nodes
for edge in state.in_edges(scope.entry):
state.add_edge(edge.src, edge.src_conn, new_entry, edge.dst_conn,
copy.deepcopy(edge.data))
for edge in state.out_edges(exit_node):
state.add_edge(new_exit, edge.src_conn, edge.dst, edge.dst_conn,
copy.deepcopy(edge.data))
# Set the exit node's map to match the entry node
new_exit.map = new_entry.map
return ScopeSubgraphView(state, new_nodes, new_entry)
def expansion(node: 'Reduce',
state: SDFGState,
sdfg: SDFG,
partial_width=16):
'''
:param node: the node to expand
:param state: the state in which the node is in
:param sdfg: the SDFG in which the node is in
:param partial_width: Width of the inner reduction buffer. Must be
larger than the latency of the reduction operation on the given
data type
'''
node.validate(sdfg, state)
inedge: graph.MultiConnectorEdge = state.in_edges(node)[0]
outedge: graph.MultiConnectorEdge = state.out_edges(node)[0]
input_dims = len(inedge.data.subset)
output_dims = len(outedge.data.subset)
input_data = sdfg.arrays[inedge.data.data]
output_data = sdfg.arrays[outedge.data.data]
# Standardize axes
axes = node.axes if node.axes else [i for i in range(input_dims)]
# Create nested SDFG
nsdfg = SDFG('reduce')
nsdfg.add_array('_in',
inedge.data.subset.size(),
input_data.dtype,
strides=input_data.strides,
storage=input_data.storage)
nsdfg.add_array('_out',
outedge.data.subset.size(),
output_data.dtype,
strides=output_data.strides,
storage=output_data.storage)
if input_data.dtype.veclen > 1:
raise NotImplementedError(
'Vectorization currently not implemented for FPGA expansion of Reduce.'
)
nstate = nsdfg.add_state()
# (If axes != all) Add outer map, which corresponds to the output range
if len(axes) != input_dims:
all_axis = False
# Interleave input and output axes to match input memlet
ictr, octr = 0, 0
input_subset = []
for i in range(input_dims):
if i in axes:
input_subset.append(f'_i{ictr}')
ictr += 1
else:
input_subset.append(f'_o{octr}')
octr += 1
output_size = outedge.data.subset.size()
ome, omx = nstate.add_map(
'reduce_output', {
f'_o{i}': f'0:{symstr(sz)}'
for i, sz in enumerate(outedge.data.subset.size())
})
outm_idx = ','.join([f'_o{i}' for i in range(output_dims)])
outm = dace.Memlet(f'_out[{outm_idx}]')
inm_idx = ','.join(input_subset)
inmm = dace.Memlet(f'_in[{inm_idx}]')
else:
all_axis = True
ome, omx = None, None
outm = dace.Memlet('_out[0]')
inm_idx = ','.join([f'_i{i}' for i in range(len(axes))])
inmm = dace.Memlet(f'_in[{inm_idx}]')
# Add inner map, which corresponds to the range to reduce
r = nstate.add_read('_in')
w = nstate.add_read('_out')
# TODO support vectorization
buffer_name = 'partial_results'
nsdfg.add_array(buffer_name, (partial_width, ),
input_data.dtype,
transient=True,
storage=dtypes.StorageType.FPGA_Local)
buffer = nstate.add_access(buffer_name)
buffer_write = nstate.add_write(buffer_name)
# Initialize explicitly partial results, as the inner map could run for a number of iteration < partial_width
init_me, init_mx = nstate.add_map(
'partial_results_init', {'i': f'0:{partial_width}'},
schedule=dtypes.ScheduleType.FPGA_Device,
unroll=True)
init_tasklet = nstate.add_tasklet('init_pr', {}, {'pr_out'},
f'pr_out = {node.identity}')
nstate.add_memlet_path(init_me, init_tasklet, memlet=dace.Memlet())
nstate.add_memlet_path(init_tasklet,
init_mx,
buffer,
src_conn='pr_out',
memlet=dace.Memlet(f'{buffer_name}[i]'))
if not all_axis:
nstate.add_memlet_path(ome, init_me, memlet=dace.Memlet())
ime, imx = nstate.add_map(
'reduce_values', {
f'_i{i}': f'0:{symstr(inedge.data.subset.size()[axis])}'
for i, axis in enumerate(sorted(axes))
})
# Accumulate over partial results
redtype = detect_reduction_type(node.wcr)
if redtype not in ExpandReduceFPGAPartialReduction._REDUCTION_TYPE_EXPR:
raise ValueError('Reduction type not supported for "%s"' %
node.wcr)
else:
reduction_expr = ExpandReduceFPGAPartialReduction._REDUCTION_TYPE_EXPR[
redtype]
# generate flatten index considering inner map: will be used for indexing into partial results
ranges_size = ime.range.size()
inner_index = '+'.join(
[f'_i{i} * {ranges_size[i + 1]}' for i in range(len(axes) - 1)])
inner_op = ' + ' if len(axes) > 1 else ''
inner_index = inner_index + f'{inner_op}_i{(len(axes) - 1)}'
partial_reduce_tasklet = nstate.add_tasklet(
'partial_reduce', {'data_in', 'buffer_in'}, {'buffer_out'}, f'''\
prev = buffer_in
buffer_out = {reduction_expr}''')
if not all_axis:
# Connect input and partial sums
nstate.add_memlet_path(r,
ome,
ime,
partial_reduce_tasklet,
dst_conn='data_in',
memlet=inmm)
else:
nstate.add_memlet_path(r,
ime,
partial_reduce_tasklet,
dst_conn='data_in',
memlet=inmm)
nstate.add_memlet_path(
buffer,
ime,
partial_reduce_tasklet,
dst_conn='buffer_in',
memlet=dace.Memlet(
f'{buffer_name}[({inner_index})%{partial_width}]'))
nstate.add_memlet_path(
partial_reduce_tasklet,
imx,
buffer_write,
src_conn='buffer_out',
memlet=dace.Memlet(
f'{buffer_name}[({inner_index})%{partial_width}]'))
# Then perform reduction on partial results
reduce_entry, reduce_exit = nstate.add_map(
'reduce', {'i': f'0:{partial_width}'},
schedule=dtypes.ScheduleType.FPGA_Device,
unroll=True)
reduce_tasklet = nstate.add_tasklet(
'reduce', {'reduce_in', 'data_in'}, {'reduce_out'}, f'''\
prev = reduce_in if i > 0 else {node.identity}
reduce_out = {reduction_expr}''')
nstate.add_memlet_path(buffer_write,
reduce_entry,
reduce_tasklet,
dst_conn='data_in',
memlet=dace.Memlet(f'{buffer_name}[i]'))
reduce_name = 'reduce_result'
nsdfg.add_array(reduce_name, (1, ),
output_data.dtype,
transient=True,
storage=dtypes.StorageType.FPGA_Local)
reduce_read = nstate.add_access(reduce_name)
reduce_access = nstate.add_access(reduce_name)
if not all_axis:
nstate.add_memlet_path(ome, reduce_read, memlet=dace.Memlet())
nstate.add_memlet_path(reduce_read,
reduce_entry,
reduce_tasklet,
dst_conn='reduce_in',
memlet=dace.Memlet(f'{reduce_name}[0]'))
nstate.add_memlet_path(reduce_tasklet,
reduce_exit,
reduce_access,
src_conn='reduce_out',
memlet=dace.Memlet(f'{reduce_name}[0]'))
if not all_axis:
# Write out the result
nstate.add_memlet_path(reduce_access, omx, w, memlet=outm)
else:
nstate.add_memlet_path(reduce_access, w, memlet=outm)
# Rename outer connectors and add to node
inedge._dst_conn = '_in'
outedge._src_conn = '_out'
node.add_in_connector('_in')
node.add_out_connector('_out')
nsdfg.validate()
return nsdfg
def can_be_applied(self,
graph: SDFGState,
expr_index,
sdfg,
permissive=False):
# Ensure both arrays contain the same data
arr1 = self.array1
arr2 = self.array2
if arr1.data != arr2.data:
return False
# Ensure only arr1's node ID contains incoming edges
if graph.in_degree(arr2) > 0:
return False
# Ensure arr1 and arr2's node IDs are ordered (avoid duplicates)
arr1_id = graph.node_id(self.array1)
arr2_id = graph.node_id(self.array2)
if (graph.in_degree(arr1) == 0 and graph.in_degree(arr2) == 0
and arr1_id >= arr2_id):
return False
map = self.map_entry
# If array's connector leads directly to map, skip it
if all(e.dst_conn and not e.dst_conn.startswith('IN_')
for e in graph.edges_between(arr1, map)):
return False
if all(e.dst_conn and not e.dst_conn.startswith('IN_')
for e in graph.edges_between(arr2, map)):
return False
if (any(e.dst != map for e in graph.out_edges(arr1))
or any(e.dst != map for e in graph.out_edges(arr2))):
return False
# Ensure arr1 and arr2 are the first two incoming nodes (avoid further
# duplicates)
all_source_nodes = set(
graph.node_id(e.src) for e in graph.in_edges(map) if e.src != arr1
and e.src != arr2 and e.src.data == arr1.data and e.dst_conn
and e.dst_conn.startswith('IN_') and graph.in_degree(e.src) == 0)
if any(nid < arr1_id or nid < arr2_id for nid in all_source_nodes):
return False
return True
def can_be_applied(graph: SDFGState,
candidate: Dict[xf.PatternNode, int],
expr_index: int,
sdfg: SDFG,
strict: bool = False) -> bool:
access = graph.node(candidate[StreamingMemory.access])
# Make sure the access node is only accessed once (read or write),
# and not at the same time
if graph.out_degree(access) > 0 and graph.in_degree(access) > 0:
return False
# If already a stream, skip
if isinstance(sdfg.arrays[access.data], data.Stream):
return False
# If does not exist on off-chip memory, skip
if sdfg.arrays[access.data].storage not in [
dtypes.StorageType.CPU_Heap, dtypes.StorageType.CPU_Pinned,
dtypes.StorageType.GPU_Global, dtypes.StorageType.FPGA_Global
]:
return False
# Only free nodes are allowed (search up the SDFG tree)
curstate = graph
node = access
while curstate is not None:
if curstate.entry_node(node) is not None:
return False
if curstate.parent.parent_nsdfg_node is None:
break
node = curstate.parent.parent_nsdfg_node
curstate = curstate.parent.parent
# Only one memlet path is allowed per outgoing/incoming edge
edges = (graph.out_edges(access)
if expr_index == 0 else graph.in_edges(access))
for edge in edges:
mpath = graph.memlet_path(edge)
if len(mpath) != len(list(graph.memlet_tree(edge))):
return False
# The innermost end of the path must have a clearly defined memory
# access pattern
innermost_edge = mpath[-1] if expr_index == 0 else mpath[0]
if (innermost_edge.data.subset.num_elements() != 1
or innermost_edge.data.dynamic
or innermost_edge.data.volume != 1):
return False
# Check if any of the maps has a dynamic range
# These cases can potentially work but some nodes (and perhaps
# tasklets) need to be replicated, which are difficult to track.
for pe in mpath:
node = pe.dst if expr_index == 0 else graph.entry_node(pe.src)
if isinstance(
node,
nodes.MapEntry) and sdutil.has_dynamic_map_inputs(
graph, node):
return False
# If already applied on this memlet and this is the I/O component, skip
if expr_index == 0:
other_node = graph.node(candidate[StreamingMemory.entry])
else:
other_node = graph.node(candidate[StreamingMemory.exit])
other_node = graph.entry_node(other_node)
if other_node.label.startswith('__s'):
return False
return True
def gemv_libnode(sdfg: SDFG,
state: SDFGState,
A,
x,
y,
alpha,
beta,
trans=None):
# Get properties
if trans is None:
trans = (sdfg.arrays[x].shape[0] == sdfg.arrays[A].shape[0])
# Add nodes
A_in, x_in = (state.add_read(name) for name in (A, x))
y_out = state.add_write(y)
libnode = Gemv('gemv', transA=trans, alpha=alpha, beta=beta)
state.add_node(libnode)
# Connect nodes
state.add_edge(A_in, None, libnode, '_A', mm.Memlet(A))
state.add_edge(x_in, None, libnode, '_x', mm.Memlet(x))
state.add_edge(libnode, '_y', y_out, None, mm.Memlet(y))
if beta != 0:
y_in = state.add_read(y)
state.add_edge(y_in, None, libnode, '_y', mm.Memlet(y))
return []
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 _create_ceil_range(self, sdfg: SDFG, graph: SDFGState,
map_entry: nodes.MapEntry):
map_exit = graph.exit_node(map_entry)
# Retrieve transformation properties.
dim_idx = self.dim_idx
new_dim_prefix = self.new_dim_prefix
tile_size = self.tile_size
divides_evenly = self.divides_evenly
strided = self.strided
offset = self.tile_offset
tile_stride = self.tile_stride
if tile_stride == 0:
tile_stride = tile_size
# Retrieve parameter and range of dimension to be strip-mined.
target_dim = map_entry.map.params[dim_idx]
td_from, td_to, td_step = map_entry.map.range[dim_idx]
# Create new map. Replace by cloning map object?
new_dim = self._find_new_dim(sdfg, graph, map_entry, new_dim_prefix,
target_dim)
nd_from = 0
if tile_stride == 1:
nd_to = td_to - td_from
else:
nd_to = symbolic.pystr_to_symbolic(
'int_ceil(%s + 1 - %s, %s) - 1' %
(symbolic.symstr(td_to), symbolic.symstr(td_from),
symbolic.symstr(tile_stride)))
nd_step = 1
new_dim_range = (nd_from, nd_to, nd_step)
new_map = nodes.Map(new_dim + '_' + map_entry.map.label, [new_dim],
subsets.Range([new_dim_range]))
# Change the range of the selected dimension to iterate over a single
# tile
if strided:
td_from_new = symbolic.pystr_to_symbolic(new_dim)
td_to_new_approx = td_to
td_step = tile_size
elif offset == 0:
td_from_new = symbolic.pystr_to_symbolic(
'%s + %s * %s' %
(symbolic.symstr(td_from), symbolic.symstr(new_dim),
symbolic.symstr(tile_stride)))
td_to_new_exact = symbolic.pystr_to_symbolic(
'min(%s + 1, %s + %s * %s + %s) - 1' %
(symbolic.symstr(td_to), symbolic.symstr(td_from),
symbolic.symstr(tile_stride), symbolic.symstr(new_dim),
symbolic.symstr(tile_size)))
td_to_new_approx = symbolic.pystr_to_symbolic(
'%s + %s * %s + %s - 1' %
(symbolic.symstr(td_from), symbolic.symstr(tile_stride),
symbolic.symstr(new_dim), symbolic.symstr(tile_size)))
else:
# include offset
td_from_new_exact = symbolic.pystr_to_symbolic(
'max(%s,%s + %s * %s - %s)' %
(symbolic.symstr(td_from), symbolic.symstr(td_from),
symbolic.symstrtr(tile_stride), symbolic.symstr(new_dim),
symbolic.symstr(offset)))
td_from_new_approx = symbolic.pystr_to_symbolic(
'%s + %s * %s - %s ' %
(symbolic.symstr(td_from), symbolic.symstr(tile_stride),
symbolic.symstr(new_dim), symbolic.symstr(offset)))
td_from_new = dace.symbolic.SymExpr(td_from_new_exact,
td_from_new_approx)
td_to_new_exact = symbolic.pystr_to_symbolic(
'min(%s + 1, %s + %s * %s + %s - %s) -1' %
(symbolic.symstr(td_to), symbolic.symstr(td_from),
symbolic.symstr(tile_stride), symbolic.symstr(new_dim),
symbolic.symstr(tile_size), symbolic.symstr(offset)))
td_to_new_approx = symbolic.pystr_to_symbolic(
'%s + %s * %s + %s - %s - 1' %
(symbolic.symstr(td_from), symbolic.symstr(tile_stride),
symbolic.symstr(new_dim), symbolic.symstr(tile_size),
symbolic.symstr(offset)))
if divides_evenly or strided:
td_to_new = td_to_new_approx
else:
td_to_new = dace.symbolic.SymExpr(td_to_new_exact,
td_to_new_approx)
return new_dim, new_map, (td_from_new, td_to_new, td_step)
def apply(self, state: SDFGState, sdfg: SDFG):
nsdfg: nodes.NestedSDFG = self.nsdfg
new_state = sdfg.add_state_before(state)
isedge = sdfg.edges_between(new_state, state)[0]
# Find relevant symbol and data descriptor mapping
mapping: Dict[str, str] = {}
mapping.update({k: str(v) for k, v in nsdfg.symbol_mapping.items()})
mapping.update({
k: next(iter(state.in_edges_by_connector(nsdfg, k))).data.data
for k in nsdfg.in_connectors
})
mapping.update({
k: next(iter(state.out_edges_by_connector(nsdfg, k))).data.data
for k in nsdfg.out_connectors
})
# Get internal state and interstate edge
source_state = nsdfg.sdfg.start_state
nisedge = nsdfg.sdfg.out_edges(source_state)[0]
# Add state contents (nodes)
new_state.add_nodes_from(source_state.nodes())
# Replace data descriptors and symbols on state graph
for node in source_state.nodes():
if isinstance(node, nodes.AccessNode) and node.data in mapping:
node.data = mapping[node.data]
for edge in source_state.edges():
edge.data.replace(mapping)
if edge.data.data in mapping:
edge.data.data = mapping[edge.data.data]
# Add state contents (edges)
for edge in source_state.edges():
new_state.add_edge(edge.src, edge.src_conn, edge.dst,
edge.dst_conn, edge.data)
# Safe replacement of edge contents
def replfunc(m):
for k, v in mapping.items():
nisedge.data.replace(k, v, replace_keys=False)
symbolic.safe_replace(mapping, replfunc)
# Add interstate edge
for akey, aval in nisedge.data.assignments.items():
# Map assignment to outer edge
if akey not in sdfg.symbols and akey not in sdfg.arrays:
newname = akey
else:
newname = nsdfg.label + '_' + akey
isedge.data.assignments[newname] = aval
# Add symbol to outer SDFG
sdfg.add_symbol(newname, nsdfg.sdfg.symbols[akey])
# Add symbol mapping to nested SDFG
nsdfg.symbol_mapping[akey] = newname
isedge.data.condition = nisedge.data.condition
# Clean nested SDFG
nsdfg.sdfg.remove_node(source_state)
# Set new starting state
nsdfg.sdfg.start_state = nsdfg.sdfg.node_id(nisedge.dst)
def apply(self, outer_state: SDFGState, sdfg: SDFG):
nsdfg_node = self.nested_sdfg
nsdfg: SDFG = nsdfg_node.sdfg
if nsdfg_node.schedule is not dtypes.ScheduleType.Default:
infer_types.set_default_schedule_and_storage_types(
nsdfg, nsdfg_node.schedule)
#######################################################
# Collect and update top-level SDFG metadata
# Global/init/exit code
for loc, code in nsdfg.global_code.items():
sdfg.append_global_code(code.code, loc)
for loc, code in nsdfg.init_code.items():
sdfg.append_init_code(code.code, loc)
for loc, code in nsdfg.exit_code.items():
sdfg.append_exit_code(code.code, loc)
# Environments
for nstate in nsdfg.nodes():
for node in nstate.nodes():
if isinstance(node, nodes.CodeNode):
node.environments |= nsdfg_node.environments
# Symbols
outer_symbols = {str(k): v for k, v in sdfg.symbols.items()}
for ise in sdfg.edges():
outer_symbols.update(ise.data.new_symbols(sdfg, outer_symbols))
# Find original source/destination edges (there is only one edge per
# connector, according to match)
inputs: Dict[str, MultiConnectorEdge] = {}
outputs: Dict[str, MultiConnectorEdge] = {}
input_set: Dict[str, str] = {}
output_set: Dict[str, str] = {}
for e in outer_state.in_edges(nsdfg_node):
inputs[e.dst_conn] = e
input_set[e.data.data] = e.dst_conn
for e in outer_state.out_edges(nsdfg_node):
outputs[e.src_conn] = e
output_set[e.data.data] = e.src_conn
# Replace symbols using invocation symbol mapping
# Two-step replacement (N -> __dacesym_N --> map[N]) to avoid clashes
symbolic.safe_replace(nsdfg_node.symbol_mapping, nsdfg.replace_dict)
#######################################################
# Collect and modify interstate edges as necessary
outer_assignments = set()
for e in sdfg.edges():
outer_assignments |= e.data.assignments.keys()
inner_assignments = set()
for e in nsdfg.edges():
inner_assignments |= e.data.assignments.keys()
allnames = set(outer_symbols.keys()) | set(sdfg.arrays.keys())
assignments_to_replace = inner_assignments & (outer_assignments
| allnames)
sym_replacements: Dict[str, str] = {}
for assign in assignments_to_replace:
newname = data.find_new_name(assign, allnames)
allnames.add(newname)
outer_symbols[newname] = nsdfg.symbols.get(assign, None)
sym_replacements[assign] = newname
nsdfg.replace_dict(sym_replacements)
#######################################################
# Collect and modify access nodes as necessary
# Access nodes that need to be reshaped
# reshapes: Set(str) = set()
# for aname, array in nsdfg.arrays.items():
# if array.transient:
# continue
# edge = None
# if aname in inputs:
# edge = inputs[aname]
# if len(array.shape) > len(edge.data.subset):
# reshapes.add(aname)
# continue
# if aname in outputs:
# edge = outputs[aname]
# if len(array.shape) > len(edge.data.subset):
# reshapes.add(aname)
# continue
# if edge is not None and not InlineMultistateSDFG._check_strides(
# array.strides, sdfg.arrays[edge.data.data].strides,
# edge.data, nsdfg_node):
# reshapes.add(aname)
# Mapping from nested transient name to top-level name
transients: Dict[str, str] = {}
# All transients become transients of the parent (if data already
# exists, find new name)
for nstate in nsdfg.nodes():
for node in nstate.nodes():
if isinstance(node, nodes.AccessNode):
datadesc = nsdfg.arrays[node.data]
if node.data not in transients and datadesc.transient:
new_name = node.data
if (new_name in sdfg.arrays
or new_name in outer_symbols
or new_name in sdfg.constants):
new_name = f'{nsdfg.label}_{node.data}'
name = sdfg.add_datadesc(new_name,
datadesc,
find_new_name=True)
transients[node.data] = name
# All transients of edges between code nodes are also added to parent
for edge in nstate.edges():
if (isinstance(edge.src, nodes.CodeNode)
and isinstance(edge.dst, nodes.CodeNode)):
if edge.data.data is not None:
datadesc = nsdfg.arrays[edge.data.data]
if edge.data.data not in transients and datadesc.transient:
new_name = edge.data.data
if (new_name in sdfg.arrays
or new_name in outer_symbols
or new_name in sdfg.constants):
new_name = f'{nsdfg.label}_{edge.data.data}'
name = sdfg.add_datadesc(new_name,
datadesc,
find_new_name=True)
transients[edge.data.data] = name
# All constants (and associated transients) become constants of the parent
for cstname, (csttype, cstval) in nsdfg.constants_prop.items():
if cstname in sdfg.constants:
if cstname in transients:
newname = transients[cstname]
else:
newname = sdfg.find_new_constant(cstname)
transients[cstname] = newname
sdfg.constants_prop[newname] = (csttype, cstval)
else:
sdfg.constants_prop[cstname] = (csttype, cstval)
#######################################################
# Replace data on inlined SDFG nodes/edges
# Replace data names with their top-level counterparts
repldict = {}
repldict.update(transients)
repldict.update({
k: v.data.data
for k, v in itertools.chain(inputs.items(), outputs.items())
})
symbolic.safe_replace(repldict,
lambda m: replace_datadesc_names(nsdfg, m),
value_as_string=True)
# Add views whenever reshapes are necessary
# for dname in reshapes:
# desc = nsdfg.arrays[dname]
# # To avoid potential confusion, rename protected __return keyword
# if dname.startswith('__return'):
# newname = f'{nsdfg.name}_ret{dname[8:]}'
# else:
# newname = dname
# newname, _ = sdfg.add_view(newname,
# desc.shape,
# desc.dtype,
# storage=desc.storage,
# strides=desc.strides,
# offset=desc.offset,
# debuginfo=desc.debuginfo,
# allow_conflicts=desc.allow_conflicts,
# total_size=desc.total_size,
# alignment=desc.alignment,
# may_alias=desc.may_alias,
# find_new_name=True)
# repldict[dname] = newname
# Add extra access nodes for out/in view nodes
# inv_reshapes = {repldict[r]: r for r in reshapes}
# for nstate in nsdfg.nodes():
# for node in nstate.nodes():
# if isinstance(node,
# nodes.AccessNode) and node.data in inv_reshapes:
# if nstate.in_degree(node) > 0 and nstate.out_degree(
# node) > 0:
# # Such a node has to be in the output set
# edge = outputs[inv_reshapes[node.data]]
# # Redirect outgoing edges through access node
# out_edges = list(nstate.out_edges(node))
# anode = nstate.add_access(edge.data.data)
# vnode = nstate.add_access(node.data)
# nstate.add_nedge(node, anode, edge.data)
# nstate.add_nedge(anode, vnode, edge.data)
# for e in out_edges:
# nstate.remove_edge(e)
# nstate.add_edge(vnode, e.src_conn, e.dst,
# e.dst_conn, e.data)
# Make unique names for states
statenames = set(s.label for s in sdfg.nodes())
for nstate in nsdfg.nodes():
if nstate.label in statenames:
newname = data.find_new_name(nstate.label, statenames)
statenames.add(newname)
nstate.set_label(newname)
#######################################################
# Add nested SDFG states into top-level SDFG
outer_start_state = sdfg.start_state
sdfg.add_nodes_from(nsdfg.nodes())
for ise in nsdfg.edges():
sdfg.add_edge(ise.src, ise.dst, ise.data)
#######################################################
# Reconnect inlined SDFG
source = nsdfg.start_state
sinks = nsdfg.sink_nodes()
# Reconnect state machine
for e in sdfg.in_edges(outer_state):
sdfg.add_edge(e.src, source, e.data)
for e in sdfg.out_edges(outer_state):
for sink in sinks:
sdfg.add_edge(sink, e.dst, e.data)
# Modify start state as necessary
if outer_start_state is outer_state:
sdfg.start_state = sdfg.node_id(source)
# TODO: Modify memlets by offsetting
# If both source and sink nodes are inputs/outputs, reconnect once
# edges_to_ignore = self._modify_access_to_access(new_incoming_edges,
# nsdfg, nstate, state,
# orig_data)
# source_to_outer = {n: e.src for n, e in new_incoming_edges.items()}
# sink_to_outer = {n: e.dst for n, e in new_outgoing_edges.items()}
# # If a source/sink node is one of the inputs/outputs, reconnect it,
# # replacing memlets in outgoing/incoming paths
# modified_edges = set()
# modified_edges |= self._modify_memlet_path(new_incoming_edges, nstate,
# state, sink_to_outer, True,
# edges_to_ignore)
# modified_edges |= self._modify_memlet_path(new_outgoing_edges, nstate,
# state, source_to_outer,
# False, edges_to_ignore)
# # Reshape: add connections to viewed data
# self._modify_reshape_data(reshapes, repldict, inputs, nstate, state,
# True)
# self._modify_reshape_data(reshapes, repldict, outputs, nstate, state,
# False)
# Modify all other internal edges pertaining to input/output nodes
# for nstate in nsdfg.nodes():
# for node in nstate.nodes():
# if isinstance(node, nodes.AccessNode):
# if node.data in input_set or node.data in output_set:
# if node.data in input_set:
# outer_edge = inputs[input_set[node.data]]
# else:
# outer_edge = outputs[output_set[node.data]]
# for edge in state.all_edges(node):
# if (edge not in modified_edges
# and edge.data.data == node.data):
# for e in state.memlet_tree(edge):
# if e.data.data == node.data:
# e._data = helpers.unsqueeze_memlet(
# e.data, outer_edge.data)
# Replace nested SDFG parents with new SDFG
for nstate in nsdfg.nodes():
nstate.parent = sdfg
for node in nstate.nodes():
if isinstance(node, nodes.NestedSDFG):
node.sdfg.parent_sdfg = sdfg
node.sdfg.parent_nsdfg_node = node
#######################################################
# Remove nested SDFG and state
sdfg.remove_node(outer_state)
return nsdfg.nodes()
