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)