def _dml_disambiguate_direction_dependent_views(sdfg: dace.SDFG):
""" Consider the following subgraph:
(A) -- y --> (n) -- x --> (C)
In dace, if B is a View node and A and C are access nodes, and y and x both have data set to A.data and
B.data respectively, the semantics of the graph depend on the order in which it is executed, i.e. reversing
the subgraph doesn't perform as expected anymore. To disambiguate this case, we set y.data to the View's
data.
"""
for n, state in sdfg.all_nodes_recursive():
if isinstance(n, nd.AccessNode) and type(n.desc(sdfg)) is dt.View:
in_edges = state.in_edges(n)
out_edges = state.out_edges(n)
if len(in_edges) == 1 and len(out_edges) == 1:
A = in_edges[0].src
y = in_edges[0].data
C = out_edges[0].dst
x = out_edges[0].data
if (isinstance(A, nd.AccessNode)
and isinstance(C, nd.AccessNode) and y.data == A.data
and x.data == C.data):
# flip the memlet
y.subset, y.other_subset = y.other_subset, y.subset
y.data = n.data
y.try_initialize(sdfg, state, in_edges[0])
def _count_views(sdfg: dace.SDFG) -> int:
num = 0
for n, _ in sdfg.all_nodes_recursive():
if (isinstance(n, nodes.AccessNode)
and isinstance(sdfg.arrays[n.data], data.View)):
num += 1
return num
def find_library_nodes(
sdfg: dace.SDFG,
lib_type: dace.sdfg.nodes.LibraryNode) -> dace.nodes.MapEntry:
""" Finds the first access node by the given data name. """
return [
n for n, _ in sdfg.all_nodes_recursive() if isinstance(n, lib_type)
]
def _instrument(sdfg: dace.SDFG, instr: dace.DataInstrumentationType, ignore: Optional[str] = None):
# Set instrumentation on all access nodes
for node, _ in sdfg.all_nodes_recursive():
if isinstance(node, nodes.AccessNode):
if ignore and ignore in node.data:
node.instrument = dace.DataInstrumentationType.No_Instrumentation
else:
node.instrument = instr
def _post_expand_trafos(sdfg: dace.SDFG):
while inline_sdfgs(sdfg) or fuse_states(sdfg):
pass
sdfg.simplify()
for node, _ in sdfg.all_nodes_recursive():
if isinstance(node, dace.nodes.MapEntry):
node.collapse = len(node.range)
def get_mapped_subsets_dicts(self, interval: Interval, section: dace.SDFG):
in_subsets = dict()
out_subsets = dict()
section_origins: Dict[str, Tuple[int, int]] = dict()
min_k_offsets: Dict[str, int] = dict()
for he in (ln for ln, _ in section.all_nodes_recursive()
if isinstance(ln, (HorizontalExecutionLibraryNode,
VerticalLoopLibraryNode))):
access_collection: AccessCollector.CartesianAccessCollection = get_access_collection(
he)
for name, offsets in access_collection.offsets().items():
off: Tuple[int, int, int]
for off in offsets:
origin = (
-off[0] - he.iteration_space.i_interval.start.offset,
-off[1] - he.iteration_space.j_interval.start.offset,
)
if name not in section_origins:
section_origins[name] = origin
if name not in min_k_offsets:
min_k_offsets[name] = off[2]
section_origins[name] = (
max(section_origins[name][0], origin[0]),
max(section_origins[name][1], origin[1]),
)
min_k_offsets[name] = min(min_k_offsets[name], off[2])
access_collection = get_access_collection(section)
for name, section_origin in section_origins.items():
vl_origin = self.origins[name]
shape = section.arrays[name].shape
dimensions = array_dimensions(section.arrays[name])
subset_strs = []
idx = iter(range(3))
if dimensions[0]:
subset_strs.append("{i:+d}:{i:+d}+({I})".format(
i=vl_origin[0] - section_origin[0],
I=shape[next(idx)],
))
if dimensions[1]:
subset_strs.append("{j:+d}:{j:+d}+({J})".format(
j=vl_origin[1] - section_origin[1], J=shape[next(idx)]))
if dimensions[2]:
subset_strs.append(
"k-({k_orig}){k:+d}:k-({k_orig}){k:+d}{K:+d}".format(
k_orig=get_axis_bound_str(vl_origin[2], "__K"),
k=min_k_offsets[name],
K=shape[next(idx)],
))
data_dims = shape[sum(dimensions):]
subset_strs.extend([f"0:{d}" for d in data_dims])
subset_str = ",".join(subset_strs)
if name in access_collection.read_fields():
in_subsets[name] = subset_str
if name in access_collection.write_fields():
out_subsets[name] = subset_str
return in_subsets, out_subsets
def is_papi_used(sdfg: dace.SDFG) -> bool:
""" Returns True if any of the SDFG elements includes PAPI counter
instrumentation. """
for node, _ in sdfg.all_nodes_recursive():
if isinstance(node, nodes.EntryNode) and node.map.instrument == dace.InstrumentationType.PAPI_Counters:
return True
if hasattr(node, 'instrument') and node.instrument == dace.InstrumentationType.PAPI_Counters:
return True
return False
def _to_device(sdfg: dace.SDFG, device: str) -> None:
"""Update sdfg in place."""
if device == "gpu":
for array in sdfg.arrays.values():
array.storage = dace.StorageType.GPU_Global
for node, _ in sdfg.all_nodes_recursive():
if isinstance(node, VerticalLoopLibraryNode):
node.implementation = "block"
node.default_storage_type = dace.StorageType.GPU_Global
node.map_schedule = dace.ScheduleType.Sequential
node.tiling_map_schedule = dace.ScheduleType.GPU_Device
for _, section in node.sections:
for array in section.arrays.values():
array.storage = dace.StorageType.GPU_Global
for node, _ in section.all_nodes_recursive():
if isinstance(node, HorizontalExecutionLibraryNode):
node.map_schedule = dace.ScheduleType.GPU_ThreadBlock
else:
for node, _ in sdfg.all_nodes_recursive():
if isinstance(node, VerticalLoopLibraryNode):
node.implementation = "block"
node.tile_sizes = [8, 8]
def validate_oir_sdfg(sdfg: dace.SDFG):
from gtc.dace.nodes import VerticalLoopLibraryNode
sdfg.validate()
is_correct_node_types = all(
isinstance(n, (dace.SDFGState, dace.nodes.AccessNode,
VerticalLoopLibraryNode))
for n, _ in sdfg.all_nodes_recursive())
is_correct_data_and_dtype = all(
isinstance(array, dace.data.Array) and typestr_to_data_type(
dace_dtype_to_typestr(array.dtype)) != DataType.INVALID
for array in sdfg.arrays.values())
if not is_correct_node_types or not is_correct_data_and_dtype:
raise ValueError("Not a valid OIR-level SDFG")
def vectorize(sdfg: dace.SDFG, par: str, ignored_conns: list = []):
input_bits = set([sdfg.arrays[a].dtype.bytes * 8 for a in sdfg.arrays])
if len(input_bits) > 1:
raise NotImplementedError('Different data type sizes as inputs')
input_bit_width = list(input_bits)[0]
sdfg.apply_strict_transformations()
# FIXME: Hardcoded for the demo machine (512 bits)
util.SVE_LEN.set(512 / input_bit_width)
for node, dfg in sdfg.all_nodes_recursive():
if isinstance(node, dace.nodes.MapEntry):
if node.params[-1] == par:
node.schedule = dace.ScheduleType.SVE_Map
for c in node.out_connectors:
edges = get_connector_edges(dfg, node, c, False)
vectorize_connector(sdfg, dfg, node, par, c, False)
for e in edges:
vectorize_connector(sdfg, dfg, e.dst, par, e.dst_conn,
True)
for edge, dfg in sdfg.all_edges_recursive():
if not isinstance(dfg, dace.SDFGState):
continue
# Force every output connector within the graph to be a vector
#if edge.data.wcr is None:
# continue
scope = util.get_sve_scope(sdfg, dfg, edge.src)
if scope is not None:
vectorize_connector(sdfg, dfg, edge.src, par, edge.src_conn, False)
# Then use a tweaked (but incorrect) version of infer_connector_types
infer_connector_types(sdfg)
return sdfg
def find_node(sdfg: dace.SDFG, label: str) -> dace.nodes.Node:
for n, _ in sdfg.all_nodes_recursive():
if n.label == label:
return n
def find_map_by_param(sdfg: dace.SDFG, pname: str) -> dace.nodes.MapEntry:
""" Finds the first map entry node by the given parameter name. """
return next(n for n, _ in sdfg.all_nodes_recursive()
if isinstance(n, dace.nodes.MapEntry) and pname in n.params)
def find_mapexit_by_param(sdfg: dace.SDFG, pname: str) -> dace.nodes.MapExit:
""" Finds the first map exit node by the given parameter name. """
state, entry = next(
(p, n) for n, p in sdfg.all_nodes_recursive()
if isinstance(n, dace.nodes.MapEntry) and pname in n.params)
return state.exit_node(entry)
def find_access_node(sdfg: dace.SDFG, name: str) -> dace.nodes.MapEntry:
""" Finds the first access node by the given data name. """
return next(n for n, _ in sdfg.all_nodes_recursive()
if isinstance(n, dace.nodes.AccessNode) and n.data == name)
def find_map_and_state_by_param(
sdfg: dace.SDFG,
pname: str) -> Tuple[dace.nodes.MapEntry, dace.SDFGState]:
""" Finds the first map entry node by the given parameter name. """
return next((n, p) for n, p in sdfg.all_nodes_recursive()
if isinstance(n, dace.nodes.MapEntry) and pname in n.params)
def canonicalize_sdfg(sdfg: dace.SDFG, symbols={}):
# Clean up unnecessary subgraphs
remove_scalar_transients(sdfg)
remove_unused_sinks(sdfg)
remove_constant_stencils(sdfg)
split_condition_interstate_edges(sdfg)
# Fuse and nest parallel K-loops
sdfg.apply_transformations_repeated(MapFission, validate=False)
standardize_data_layout(sdfg)
sdfg.apply_transformations_repeated([NestK, InlineSDFG], validate=False)
sdfg.apply_transformations_repeated([StencilFusion])
# Remove loops
loops_removed = sdfg.apply_transformations_repeated([RemoveLoop],
validate=False)
if loops_removed > 0:
raise ValueError("Control flow loops not supported.")
from dace.transformation.interstate import StateFusion
sdfg.apply_transformations_repeated(StateFusion)
sdfg.apply_strict_transformations()
# Specialize symbols and constants
sdfg.specialize(symbols)
symbols.update(sdfg.constants)
undefined_symbols = sdfg.free_symbols
if len(undefined_symbols) != 0:
raise ValueError("Missing symbols: {}".format(
", ".join(undefined_symbols)))
for node, _ in sdfg.all_nodes_recursive():
if isinstance(node, stencil.Stencil):
node.shape = _specialize_symbols(node.shape, symbols)
if isinstance(node, dace.sdfg.nodes.MapEntry):
ranges = []
for r in node.map.range:
ranges.append(_specialize_symbols(r, symbols))
node.map.range = ranges
# Make transformation passes on tasklets and stencil libnodes
if hasattr(node, 'code'):
new_code = [_Predicator().visit(stmt) for stmt in node.code.code]
# min/max predication requires multiple passes (nested expressions)
minmax_predicated = 1
while minmax_predicated > 0:
pred = _MinMaxPredicator()
tmp_code = [pred.visit(stmt) for stmt in new_code]
minmax_predicated = pred.count
# Some of the outputs may be lists, flatten
new_code = []
def flatten(val):
for v in val:
if isinstance(v, list):
flatten(v)
else:
new_code.append(v)
flatten(tmp_code)
node.code.code = new_code
return sdfg
def find_map(sdfg: dace.SDFG, condition=None) -> dace.nodes.MapEntry:
""" Finds the first map entry node by the given parameter name. """
return next(n for n, s in sdfg.all_nodes_recursive()
if isinstance(n, dace.nodes.MapEntry) and (
condition is None or condition(n, s)))
def iter_vertical_loop_section_sub_sdfgs(graph: SDFG) -> Iterator[SDFG]:
from gtc.dace.nodes import VerticalLoopLibraryNode
for node, _ in graph.all_nodes_recursive():
if isinstance(node, VerticalLoopLibraryNode):
yield from (subgraph for _, subgraph in node.sections)
