diff --git a/conftest.py b/conftest.py new file mode 100644 index 0000000000..669e2cc9f8 --- /dev/null +++ b/conftest.py @@ -0,0 +1,74 @@ +# Copyright 2019-2022 ETH Zurich and the DaCe authors. All rights reserved. +""" +Root pytest configuration file. +""" +import os +import subprocess + +import pytest + + +def parse_worker_index(worker_id: str) -> int: + """Parse the digits in a worker id or MPI rank string (e.g. 'gw3' -> 3, '12' -> 12). + + Falls back to 0 for text with no digits (e.g. the non-xdist 'master' worker id). + """ + digits = ''.join(char for char in worker_id if char.isdigit()) + return int(digits) if digits else 0 + + +def list_cuda_devices() -> list: + """Enumerate CUDA device indices via `nvidia-smi -L`, without touching CUDA in this process.""" + try: + result = subprocess.run(['nvidia-smi', '-L'], capture_output=True, text=True, timeout=10) + except (OSError, subprocess.TimeoutExpired): + return [] + if result.returncode != 0: + return [] + device_lines = [line for line in result.stdout.splitlines() if line.startswith('GPU ')] + return [str(index) for index in range(len(device_lines))] + + +def pick_gpu_worker_device(worker_id: str, device_pool: list) -> str: + """Round-robin a worker id (xdist worker or MPI rank) onto one device of device_pool.""" + index = parse_worker_index(worker_id) + return device_pool[index % len(device_pool)] + + +def resolve_worker_id() -> str: + """Id that places this process in a device rotation. + + The pytest-xdist worker id if present, else the first set MPI/job-launcher rank env var + (dace.sdfg.sdfg.LAUNCHER_RANK_VARS, the same list dace itself reads for per-rank build + folders), else '' if this process is neither an xdist worker nor an MPI rank. + """ + xdist_worker = os.environ.get('PYTEST_XDIST_WORKER') + if xdist_worker is not None: + return xdist_worker + from dace.sdfg.sdfg import LAUNCHER_RANK_VARS + for var in LAUNCHER_RANK_VARS: + rank = os.environ.get(var) + if rank: + return rank + return '' + + +def pin_worker_to_gpu() -> None: + """Pin this worker process (pytest-xdist or an MPI-launched rank) to a single GPU. + + Without this, every worker/rank sees all GPUs and piles its CUDA context onto device 0, which + is the root cause of flaky 'invalid device ordinal' failures under high worker/rank counts. + """ + worker_id = resolve_worker_id() + if not worker_id: + return + preset = os.environ.get('CUDA_VISIBLE_DEVICES', '') + device_pool = [entry.strip() for entry in preset.split(',') if entry.strip()] if preset.strip() \ + else list_cuda_devices() + if len(device_pool) <= 1: + return + os.environ['CUDA_VISIBLE_DEVICES'] = pick_gpu_worker_device(worker_id, device_pool) + + +def pytest_configure(config: pytest.Config) -> None: + pin_worker_to_gpu() diff --git a/dace/frontend/python/astutils.py b/dace/frontend/python/astutils.py index ec192d15d0..7e25abd2c9 100644 --- a/dace/frontend/python/astutils.py +++ b/dace/frontend/python/astutils.py @@ -1,4 +1,4 @@ -# Copyright 2019-2025 ETH Zurich and the DaCe authors. All rights reserved. +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. """ Various AST parsing utilities for DaCe. """ import ast import astunparse @@ -216,6 +216,16 @@ def _Constant(self, t): else: super()._Constant(t) + def _Attribute(self, t): + self.dispatch(t.value) + # Special case: 3.__abs__() is a syntax error, so if t.value is an integer literal + # then we need to add an extra space to get 3 .__abs__(). astunparse checks this via + # ``ast.Num``, which was removed in Python 3.12; an int Constant is the same check. + if isinstance(t.value, ast.Constant) and isinstance(t.value.value, int): + self.write(" ") + self.write(".") + self.write(t.attr) + def _Subscript(self, t): self.dispatch(t.value) self.write('[') diff --git a/dace/sdfg/performance_evaluation/assumptions.py b/dace/sdfg/performance_evaluation/assumptions.py index 1b1d37348b..5ca5d083bb 100644 --- a/dace/sdfg/performance_evaluation/assumptions.py +++ b/dace/sdfg/performance_evaluation/assumptions.py @@ -1,8 +1,10 @@ -# Copyright 2019-2023 ETH Zurich and the DaCe authors. All rights reserved. +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. import sympy as sp from typing import Dict +from dace.symbolic import symbol + class UnionFind: """ @@ -112,7 +114,9 @@ def propagate_assumptions(x, y, condensed_assumptions): condensed_assumptions[y] = Assumptions() assum_y = condensed_assumptions[y] for e in assum_x.equal: - if e is not sp.Symbol(y): + # Skip a self-reference (the symbol named ``y`` itself); compare by name because DaCe + # symbols, unlike sympy symbols, are not interned, so an ``is`` check would never match. + if not (isinstance(e, sp.Symbol) and e.name == y): assum_y.add_equal(e) for g in assum_x.greater: assum_y.add_greater(g) @@ -150,7 +154,7 @@ def propagate_assumptions_equal_symbols(condensed_assumptions): for other in condensed_assumptions[sym].equal: if isinstance(other, sp.Symbol): propagate_assumptions(sym, uf.find(sym), condensed_assumptions) - equality_subs1.update({sym: sp.Symbol(uf.find(sym))}) + equality_subs1.update({sym: symbol(uf.find(sym))}) equality_subs2 = {} # In a second step, each symbol gets replaced with its equal number (if present) @@ -221,37 +225,37 @@ def parse_assumptions(assumptions, array_symbols): condensed_assumptions: Dict[str, Assumptions] = {} for a in assumptions: if '==' in a: - symbol, rhs = a.split('==') - if symbol not in condensed_assumptions: - condensed_assumptions[symbol] = Assumptions() + lhs, rhs = a.split('==') + if lhs not in condensed_assumptions: + condensed_assumptions[lhs] = Assumptions() try: - condensed_assumptions[symbol].add_equal(int(rhs)) + condensed_assumptions[lhs].add_equal(int(rhs)) except ValueError: - condensed_assumptions[symbol].add_equal(sp.Symbol(rhs)) + condensed_assumptions[lhs].add_equal(symbol(rhs)) elif '>' in a: - symbol, rhs = a.split('>') - if symbol not in condensed_assumptions: - condensed_assumptions[symbol] = Assumptions() + lhs, rhs = a.split('>') + if lhs not in condensed_assumptions: + condensed_assumptions[lhs] = Assumptions() try: - condensed_assumptions[symbol].add_greater(int(rhs)) + condensed_assumptions[lhs].add_greater(int(rhs)) except ValueError: - condensed_assumptions[symbol].add_greater(sp.Symbol(rhs)) + condensed_assumptions[lhs].add_greater(symbol(rhs)) # add the opposite, i.e. for x>y, we add yx if rhs not in condensed_assumptions: condensed_assumptions[rhs] = Assumptions() - condensed_assumptions[rhs].add_greater(sp.Symbol(symbol)) + condensed_assumptions[rhs].add_greater(symbol(lhs)) # Handle equal assumptions. equality_subs = propagate_assumptions_equal_symbols(condensed_assumptions) @@ -272,14 +276,14 @@ def parse_assumptions(assumptions, array_symbols): for sym, assum in condensed_assumptions.items(): i = 0 for g in assum.greater: - replacement_symbol = sp.Symbol(f'_p_{sym}', positive=True, integer=True) - all_subs[i][0].update({sp.Symbol(sym): replacement_symbol + g}) - all_subs[i][1].update({replacement_symbol: sp.Symbol(sym) - g}) + replacement_symbol = symbol(f'_p_{sym}', nonnegative=True) + all_subs[i][0].update({symbol(sym): replacement_symbol + g}) + all_subs[i][1].update({replacement_symbol: symbol(sym) - g}) i += 1 for l in assum.lesser: - replacement_symbol = sp.Symbol(f'_n_{sym}', negative=True, integer=True) - all_subs[i][0].update({sp.Symbol(sym): replacement_symbol + l}) - all_subs[i][1].update({replacement_symbol: sp.Symbol(sym) - l}) + replacement_symbol = symbol(f'_n_{sym}', negative=True) + all_subs[i][0].update({symbol(sym): replacement_symbol + l}) + all_subs[i][1].update({replacement_symbol: symbol(sym) - l}) i += 1 return equality_subs, all_subs diff --git a/dace/sdfg/performance_evaluation/helpers.py b/dace/sdfg/performance_evaluation/helpers.py index ba7bfb84f2..ab4028d5cf 100644 --- a/dace/sdfg/performance_evaluation/helpers.py +++ b/dace/sdfg/performance_evaluation/helpers.py @@ -1,26 +1,15 @@ -# Copyright 2019-2023 ETH Zurich and the DaCe authors. All rights reserved. -""" Helper functions used by the work depth analysis. """ +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. +""" Helper functions shared by the SDFG performance analyses: element UUIDs, fixed-point symbol +substitution, and static-symbol detection. """ -from dace import SDFG, SDFGState, nodes -from collections import deque -from typing import List, Dict, Set, Tuple, Optional, Union -import networkx as nx +import re +from typing import Dict -NodeT = str -EdgeT = Tuple[NodeT, NodeT] - - -class NodeCycle: - - nodes: Set[NodeT] = [] - - def __init__(self, nodes: List[NodeT]) -> None: - self.nodes = set(nodes) - - @property - def length(self) -> int: - return len(self.nodes) +import sympy as sp +from dace import SDFG, SDFGState, dtypes, nodes +from dace.sdfg.state import BreakBlock, ContinueBlock, ReturnBlock, UnstructuredControlFlow +from dace.symbolic import pystr_to_symbolic, symbol UUID_SEPARATOR = '/' @@ -41,297 +30,122 @@ def get_uuid(element, state=None): return ids_to_string(-1) -def get_domtree(graph: nx.DiGraph, start_node: str, idom: Dict[str, str] = None): - idom = idom or nx.immediate_dominators(graph, start_node) - - alldominated = {n: set() for n in graph.nodes} - domtree = nx.DiGraph() - - for node, dom in idom.items(): - if node is dom: - continue - domtree.add_edge(dom, node) - alldominated[dom].add(node) - - nextidom = idom[dom] - ndom = nextidom if nextidom != dom else None - - while ndom: - alldominated[ndom].add(node) - nextidom = idom[ndom] - ndom = nextidom if nextidom != ndom else None - - # 'Rank' the tree, i.e., annotate each node with the level it is on. - q = deque() - q.append((start_node, 0)) - while q: - node, level = q.popleft() - domtree.add_node(node, level=level) - for s in domtree.successors(node): - q.append((s, level + 1)) - - return alldominated, domtree - - -def get_backedges(graph: nx.DiGraph, - start: Optional[NodeT], - strict: bool = False) -> Union[Set[EdgeT], Tuple[Set[EdgeT], Set[EdgeT]]]: - '''Find all backedges in a directed graph. - - Note: - This algorithm has an algorithmic complexity of O((|V|+|E|)*C) for a - graph with vertices V, edges E, and C cycles. - - Args: - graph (nx.DiGraph): The graph for which to search backedges. - start (str): Start node of the graph. If no start is provided, a node - with no incoming edges is used as the start. If no such node can - be found, a `ValueError` is raised. - - Returns: - A set of backedges in the graph. - - Raises: - ValueError: If no `start` is provided and the graph contains no nodes - with no incoming edges. - ''' - backedges = set() - eclipsed_backedges = set() - - if start is None: - for node in graph.nodes(): - if graph.in_degree(node) == 0: - start = node - break - if start is None: - raise ValueError('No start node provided and no start node could ' + 'be determined automatically') - - # Gather all cycles in the graph. Cycles are represented as a sequence of - # nodes. - # O((|V|+|E|)*(C+1)), for C cycles. - all_cycles_nx: List[List[NodeT]] = nx.cycles.simple_cycles(graph) - #all_cycles_nx: List[List[NodeT]] = nx.simple_cycles(graph) - all_cycles: Set[NodeCycle] = set() - for cycle in all_cycles_nx: - all_cycles.add(NodeCycle(cycle)) - - # Construct a dictionary mapping a node to the cycles containing that node. - # O(|V|*|C|) - cycle_map: Dict[NodeT, Set[NodeCycle]] = dict() - for cycle in all_cycles: - for node in cycle.nodes: - try: - cycle_map[node].add(cycle) - except KeyError: - cycle_map[node] = set([cycle]) - - # Do a BFS traversal of the graph to detect the back edges. - # For each node that is part of an (unhandled) cycle, find the longest - # still unhandled cycle and try to use it to find the back edge for it. - bfs_frontier = [start] - visited: Set[NodeT] = set([start]) - handled_cycles: Set[NodeCycle] = set() - unhandled_cycles = all_cycles - while bfs_frontier: - node = bfs_frontier.pop(0) - pred = [p for p in graph.predecessors(node) if p not in visited] - longest_cycles: Dict[NodeT, NodeCycle] = dict() - try: - cycles = cycle_map[node] - remove_cycles = set() - for cycle in cycles: - if cycle not in handled_cycles: - for p in pred: - if p in cycle.nodes: - if p not in longest_cycles: - longest_cycles[p] = cycle - else: - if cycle.length > longest_cycles[p].length: - longest_cycles[p] = cycle - else: - remove_cycles.add(cycle) - for cycle in remove_cycles: - cycles.remove(cycle) - except KeyError: - longest_cycles = dict() - - # For the current node, find the incoming edge which belongs to the - # cycle and has not been visited yet, which indicates a backedge. - node_backedge_candidates: Set[Tuple[EdgeT, NodeCycle]] = set() - for p, longest_cycle in longest_cycles.items(): - handled_cycles.add(longest_cycle) - unhandled_cycles.remove(longest_cycle) - cycle_map[node].remove(longest_cycle) - backedge_candidates = graph.in_edges(node) - for candidate in backedge_candidates: - src = candidate[0] - dst = candidate[0] - if src not in visited and src in longest_cycle.nodes: - node_backedge_candidates.add((candidate, longest_cycle)) - if not strict: - backedges.add(candidate) - - # Make sure that any cycle containing this back edge is - # not evaluated again, i.e., mark as handled. - remove_cycles = set() - for cycle in unhandled_cycles: - if src in cycle.nodes and dst in cycle.nodes: - handled_cycles.add(cycle) - remove_cycles.add(cycle) - for cycle in remove_cycles: - unhandled_cycles.remove(cycle) - - # If strict is set, we only report the longest cycle's back edges for - # any given node, and separately return any other backedges as - # 'eclipsed' backedges. In the case of a while-loop, for example, - # the loop edge is considered a backedge, while a continue inside the - # loop is considered an 'eclipsed' backedge. - if strict: - longest_candidate: Tuple[EdgeT, NodeCycle] = None - eclipsed_candidates = set() - for be_candidate in node_backedge_candidates: - if longest_candidate is None: - longest_candidate = be_candidate - elif longest_candidate[1].length < be_candidate[1].length: - eclipsed_candidates.add(longest_candidate[0]) - longest_candidate = be_candidate - else: - eclipsed_candidates.add(be_candidate[0]) - if longest_candidate is not None: - backedges.add(longest_candidate[0]) - if eclipsed_candidates: - eclipsed_backedges.update(eclipsed_candidates) - - # Continue BFS. - for neighbour in graph.successors(node): - if neighbour not in visited: - visited.add(neighbour) - bfs_frontier.append(neighbour) - - if strict: - return backedges, eclipsed_backedges - else: - return backedges - - -class LoopExtractionError(Exception): - pass - - -def find_loop_guards_tails_exits(sdfg_nx: nx.DiGraph): - """ - Detects loops in a SDFG. For each loop, it identifies (node, oNode, exit). - We know that there is a backedge from oNode to node that creates the loop and that exit is the exit state of the loop. - - :param sdfg_nx: The networkx representation of a SDFG. +def has_unstructured_control_flow(sdfg: SDFG) -> bool: """ + Check whether the SDFG contains control flow the performance analyses do not model. - # preparation phase: compute dominators, backedges etc - for node in sdfg_nx.nodes(): - if sdfg_nx.in_degree(node) == 0: - start = node - break - if start is None: - raise ValueError('No start node could be determined') + They assume structured control flow -- loops as ``LoopRegion`` and branches as + ``ConditionalBlock`` -- and model neither non-local exits nor legacy state machines. The + following are therefore reported as unsupported: a legacy loop (a cycle not encapsulated in a + ``LoopRegion``), unstructured branching (a block with more than one outgoing edge), and + ``break`` / ``continue`` / ``return`` (``BreakBlock`` / ``ContinueBlock`` / ``ReturnBlock``). - # sdfg can have multiple end nodes --> not good for postDomTree - # --> add a new end node - artificial_end_node = 'artificial_end_node' - sdfg_nx.add_node(artificial_end_node) - for node in sdfg_nx.nodes(): - if sdfg_nx.out_degree(node) == 0 and node != artificial_end_node: - # this is an end node of the sdfg - sdfg_nx.add_edge(node, artificial_end_node) - - # sanity check: - if sdfg_nx.in_degree(artificial_end_node) == 0: - raise LoopExtractionError('No end node could be determined in the SDFG') - - # compute dominators and backedges - iDoms = nx.immediate_dominators(sdfg_nx, start) - allDom, _ = get_domtree(sdfg_nx, start, iDoms) - - reversed_sdfg_nx = sdfg_nx.reverse() - iPostDoms = nx.immediate_dominators(reversed_sdfg_nx, artificial_end_node) - _, postDomTree = get_domtree(reversed_sdfg_nx, artificial_end_node, iPostDoms) + :param sdfg: The SDFG to inspect. + :return: True if any (possibly nested) control-flow region is unstructured. + """ + for region in sdfg.all_control_flow_regions(recursive=True): + if isinstance(region, UnstructuredControlFlow) or region.has_cycles(): + return True + for block in region.nodes(): + if isinstance(block, (BreakBlock, ContinueBlock, ReturnBlock)) or len(region.out_edges(block)) > 1: + return True + return False - backedges = get_backedges(sdfg_nx, start) - backedgesDstDict = {} - for be in backedges: - if be[1] in backedgesDstDict: - backedgesDstDict[be[1]].add(be) - else: - backedgesDstDict[be[1]] = set([be]) - # This list will be filled with triples (node, oNode, exit), one triple for each loop construct in the SDFG. - # There will always be a backedge from oNode to node. Either node or oNode will be the corresponding loop guard, - # depending on whether it is a while-do or a do-while loop. exit will always be the exit state of the loop. - nodes_oNodes_exits = [] +def subs_till_fixed_point(expr: sp.Expr, symbol_map: Dict[sp.Expr, sp.Expr]) -> sp.Expr: + """ + Apply a symbol mapping to a symbolic expression repeatedly until a fixed point is reached. - # iterate over all nodes - for node in sdfg_nx.nodes(): - # Check if any backedge ends in node. - if node in backedgesDstDict: - inc_backedges = backedgesDstDict[node] + Requires that the symbol mapping has no cyclic dependencies, otherwise it would not converge. - # gather all successors of node that are not reached by backedges - successors = [] - for edge in sdfg_nx.out_edges(node): - if not edge in backedges: - successors.append(edge[1]) + :param expr: The expression to substitute into (non-symbolic values are returned unchanged). + :param symbol_map: Mapping from symbols to their replacement expressions. + :return: The expression after substituting to a fixed point. + """ + if not isinstance(expr, sp.Expr): + return expr + prev = None + curr = expr + while prev != curr: + prev = curr + curr = curr.subs(symbol_map) + return curr - # For each incoming backedge, we want to find oNode and exit. There can be multiple backedges, in case - # we have a continue statement in the original code. But we can handle these backedges normally. - for be in inc_backedges: - # since node has an incoming backedge, it is either a loop guard or loop tail - # oNode will exactly be the other thing - oNode = be[0] - exitCandidates = set() - # search for exit candidates: - # a state is a exit candidate if: - # - it is in successor and it does not dominate oNode (else it dominates - # the last loop state, and hence is inside the loop itself) - # - is is a successor of oNode (but not node) - # This handles both cases of while-do and do-while loops - for succ in successors: - if succ != oNode and oNode not in allDom[succ]: - exitCandidates.add(succ) - for succ in sdfg_nx.successors(oNode): - if succ != node: - exitCandidates.add(succ) - if len(exitCandidates) == 0: - raise LoopExtractionError('failed to find any exit nodes') - elif len(exitCandidates) > 1: - # Find the exit candidate that sits highest up in the - # postdominator tree (i.e., has the lowest level). - # That must be the exit node (it must post-dominate) - # everything inside the loop. If there are multiple - # candidates on the lowest level (i.e., disjoint set of - # postdominated nodes), there are multiple exit paths, - # and they all share one level. - cand = exitCandidates.pop() - minSet = set([cand]) - minLevel = nx.get_node_attributes(postDomTree, 'level')[cand] - for cand in exitCandidates: - curr_level = nx.get_node_attributes(postDomTree, 'level')[cand] - if curr_level < minLevel: - # new minimum found - minLevel = curr_level - minSet.clear() - minSet.add(cand) - elif curr_level == minLevel: - # add cand to curr set - minSet.add(cand) +def get_static_symbols(sdfg: SDFG) -> Dict[str, sp.Expr]: + """ + Find the symbols that are assigned at exactly one point in the SDFG (i.e., statically known). - if len(minSet) > 0: - exitCandidates = minSet - else: - raise LoopExtractionError('failed to find exit minSet') + A symbol is static if it is written by a single length-1 access (from a tasklet performing one + assignment, or by a single copy from another access node). Symbols written in more than one + place are excluded. - # now we have a triple (node, oNode, exitCandidates) - nodes_oNodes_exits.append((node, oNode, exitCandidates)) + :param sdfg: The SDFG for which to find static symbols and their assignments. + :return: Mapping from each static symbol name to its defining expression (resolved to a fixed + point). String keys let callers both substitute and index the result by symbol name. + """ + # Strip type-cast prefixes (e.g. ``dace.float64``, ``int``) from a tasklet RHS so the cast does + # not interfere with the symbolic parse below. The DaCe type names are derived from + # ``dace.dtypes`` (rather than hard-coded), and matched longest-first so ``dace.float64`` wins + # over ``float``. + cast_names = {'int', 'float', 'complex', 'bool'} + cast_names |= {f'dace.{name}' for name in dir(dtypes) if isinstance(getattr(dtypes, name), dtypes.typeclass)} + type_regex = re.compile("|".join(re.escape(name) for name in sorted(cast_names, key=len, reverse=True))) + + static_symbol_mapping: Dict[sp.Symbol, sp.Expr] = {symbol(a): symbol(a) for a in sdfg.arg_names} + non_static_symbols = set() + for node, containing_state in sdfg.all_nodes_recursive(): + if not isinstance(node, nodes.AccessNode): + continue + if containing_state.in_degree(node) != 1: + continue + edge = containing_state.in_edges(node)[0] + source = edge.src + if edge.data.volume != 1: + continue - # remove artificial end node - sdfg_nx.remove_node(artificial_end_node) - return nodes_oNodes_exits + if isinstance(source, nodes.Tasklet): + tasklet = source + in_map = {} + out_map = {} + # Incoming edges: symbols feeding the tasklet. + for e in containing_state.in_edges(tasklet): + if not isinstance(e.src, nodes.AccessNode): + continue + in_map[e.dst_conn] = str(e.src.data) + # Outgoing edges: symbols written by the tasklet (expected to be a single edge). + for e in containing_state.out_edges(tasklet): + if not isinstance(e.dst, nodes.AccessNode): + continue + out_map[e.src_conn] = str(e.dst.data) + + in_map = {symbol(k): symbol(v) for k, v in in_map.items()} + out_map = {symbol(k): symbol(v) for k, v in out_map.items()} + code = tasklet.code.as_string.strip() + # Expect a single assignment. + lines = [l.strip() for l in code.splitlines() if l.strip()] + if len(lines) > 1: + non_static_symbols.add(node.data) + continue + lhs, rhs = lines[0].split('=', 1) + lhs = lhs.strip() + rhs = type_regex.sub("", rhs.strip()) + lhs_sympy = pystr_to_symbolic(lhs).subs(out_map) + + if lhs_sympy not in static_symbol_mapping.keys(): + try: + static_symbol_mapping[lhs_sympy] = pystr_to_symbolic(rhs).subs(in_map) + except Exception: + non_static_symbols.add(lhs_sympy) + else: + non_static_symbols.add(lhs_sympy) + + elif isinstance(source, nodes.AccessNode): + data_sym = symbol(source.data) + if data_sym not in static_symbol_mapping.keys(): + static_symbol_mapping[data_sym] = symbol(node.data) + else: + non_static_symbols.add(data_sym) + + static_symbol_mapping = {k: v for (k, v) in static_symbol_mapping.items() if k not in non_static_symbols} + return {str(k): subs_till_fixed_point(v, static_symbol_mapping) for k, v in static_symbol_mapping.items()} diff --git a/dace/sdfg/performance_evaluation/op_in_helpers.py b/dace/sdfg/performance_evaluation/op_in_helpers.py index 2375268625..ca8fe0c246 100644 --- a/dace/sdfg/performance_evaluation/op_in_helpers.py +++ b/dace/sdfg/performance_evaluation/op_in_helpers.py @@ -1,15 +1,14 @@ -# Copyright 2019-2023 ETH Zurich and the DaCe authors. All rights reserved. +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. """ Contains class CacheLineTracker which keeps track of all arrays of an SDFG and their cache line position and class AccessStack which which corresponds to the stack used to compute the stack distance. Further, provides a curve fitting method and plotting function. """ - import warnings from dace.data import Array import sympy as sp -from collections import deque from scipy.optimize import curve_fit import numpy as np from dace import symbol, symbolic +from dace.symbolic import symbol, pystr_to_symbolic class CacheLineTracker: @@ -34,18 +33,12 @@ def cache_line_id(self, name: str, access: [int], mapping): one_d_index = 0 for dim in range(len(access)): i = access[dim] - one_d_index += (i + sp.sympify(arr.offset[dim]).subs(mapping)) * sp.sympify(arr.strides[dim]).subs(mapping) + one_d_index += (i + pystr_to_symbolic(arr.offset[dim]).subs(mapping)) * pystr_to_symbolic( + arr.strides[dim]).subs(mapping) # divide by L to get the cache line id return self.start_lines[name] + symbolic.int_floor(one_d_index * arr.dtype.bytes, self.L) - def copy(self): - new_clt = CacheLineTracker(self.L) - new_clt.array_info = dict(self.array_info) - new_clt.start_lines = dict(self.start_lines) - new_clt.next_free_line = self.next_free_line - return new_clt - class Node: @@ -96,48 +89,19 @@ def touch(self, id): return distance - def in_cache_as_list(self): - """ - Returns a list of cache ids currently in cache. Index 0 is the most recently used. - """ - res = deque() - curr = self.top - dist = 0 - while curr is not None and dist < self.C: - res.append(curr.v) - curr = curr.next - dist += 1 - return res - - def debug_print(self): - # prints the whole stack - print('\n') - curr = self.top - while curr is not None: - print(curr.v, end=', ') - curr = curr.next - print('\n') - - def copy(self): - new_stack = AccessStack(self.C) - cache_content = self.in_cache_as_list() - if len(cache_content) > 0: - new_top_value = cache_content.popleft() - new_stack.top = Node(new_top_value) - curr = new_stack.top - for x in cache_content: - curr.next = Node(x) - curr = curr.next - return new_stack + def replace_self(self, other: 'AccessStack'): + self.top = other.top + self.num_calls = other.num_calls + self.lengh = other.length + self.C = other.C def plot(x, work_map, cache_misses, op_in_map, symbol_name, C, L, sympy_f, element, name): - plt = None + # matplotlib is an optional dependency, imported lazily so the rest of the module works without it. try: - import matplotlib.pyplot as plt_import - plt = plt_import + import matplotlib.pyplot as plt except ModuleNotFoundError: - pass + plt = None if plt is None: warnings.warn('Plotting only possible with matplotlib installed') @@ -154,7 +118,7 @@ def plot(x, work_map, cache_misses, op_in_map, symbol_name, C, L, sympy_f, eleme ax[0].scatter(x, cache_misses, label=f'C={C*L}, L={L}') b = [] for curr in a: - b.append(sp.N(sp.sympify(sympy_f).subs(symbol_name, curr))) + b.append(sp.N(pystr_to_symbolic(sympy_f).subs(symbol_name, curr))) ax[0].plot(a, b) c = [] @@ -170,7 +134,7 @@ def plot(x, work_map, cache_misses, op_in_map, symbol_name, C, L, sympy_f, eleme ax[1].scatter(x, c, label=f'C={C*L}, L={L}') b = [] for curr in a: - b.append(sp.N(sp.sympify(op_in_map).subs(symbol_name, curr))) + b.append(sp.N(pystr_to_symbolic(op_in_map).subs(symbol_name, curr))) ax[1].plot(a, b) ax[0].set_ylim(bottom=0, top=max(cache_misses) + max(cache_misses) / 10) diff --git a/dace/sdfg/performance_evaluation/operational_intensity.py b/dace/sdfg/performance_evaluation/operational_intensity.py index ee9286a7c5..1a4bfaa845 100644 --- a/dace/sdfg/performance_evaluation/operational_intensity.py +++ b/dace/sdfg/performance_evaluation/operational_intensity.py @@ -1,18 +1,21 @@ -# Copyright 2019-2023 ETH Zurich and the DaCe authors. All rights reserved. +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. """ Analyses the operational intensity of an input SDFG. Can be used as a Python script or from the VS Code extension. """ import argparse -from collections import deque from dace.sdfg import nodes as nd -from dace import SDFG, SDFGState, dtypes +from dace import dtypes, SDFG +from dace.sdfg.state import SDFGState, ControlFlowRegion, LoopRegion, FunctionCallRegion, ConditionalBlock from typing import Tuple, Dict import os import sympy as sp from copy import deepcopy -from dace.symbolic import pystr_to_symbolic, SymExpr +from dace.symbolic import pystr_to_symbolic, SymExpr, symbol, simplify +import re +import warnings -from dace.sdfg.performance_evaluation.helpers import get_uuid +from dace.sdfg.performance_evaluation.helpers import (get_uuid, get_static_symbols, subs_till_fixed_point, + has_unstructured_control_flow) from dace.transformation.passes.symbol_ssa import StrictSymbolSSA from dace.transformation.pass_pipeline import FixedPointPipeline @@ -20,6 +23,8 @@ from dace.sdfg.performance_evaluation.op_in_helpers import CacheLineTracker, AccessStack, fit_curve, plot, compute_mape from dace.sdfg.performance_evaluation.work_depth import analyze_sdfg, get_tasklet_work +from dace.transformation.passes.analysis import loop_analysis + class SymbolRange(): """ Used to describe an SDFG symbol associated with a range (start, stop, step) of values. """ @@ -60,7 +65,7 @@ def update_map(op_in_map, uuid, new_misses, average=True): def calculate_op_in(op_in_map, work_map, stringify=False, assumptions={}): """ Calculates the operational intensity for each SDFG element from work and bytes loaded. """ for uuid in op_in_map: - work = work_map[uuid][0].subs(assumptions) + work = work_map[uuid].subs(assumptions) if work == 0 and op_in_map[uuid] == 0: op_in_map[uuid] = 0 elif work != 0 and op_in_map[uuid] == 0: @@ -73,43 +78,6 @@ def calculate_op_in(op_in_map, work_map, stringify=False, assumptions={}): op_in_map[uuid] = str(op_in_map[uuid]) -def mem_accesses_on_path(states): - mem_accesses = 0 - for state in states: - mem_accesses += len(state.read_and_write_sets()) - return mem_accesses - - -def find_states_between(sdfg: SDFG, start_state: SDFGState, end_state: SDFGState): - traversal_q = deque() - traversal_q.append(start_state) - visited = set() - states = [] - while traversal_q: - curr_state = traversal_q.popleft() - if curr_state == end_state: - continue - if curr_state not in visited: - visited.add(curr_state) - states.append(curr_state) - for e in sdfg.out_edges(curr_state): - traversal_q.append(e.dst) - return states - - -def find_merge_state(sdfg: SDFG, state: SDFGState): - """ - Adapted from ``cfg.stateorder_topological_sort``. - """ - from dace.sdfg.analysis import cfg - - merges = cfg.branch_merges(sdfg) - if state in merges: - return merges[state] - - print(f'WARNING: No merge state could be detected for branch state "{state.name}".', ) - - def symeval(val, symbols): """ Takes a sympy expression and substitutes its symbols according to a dict { old_symbol: new_symbol}. @@ -119,7 +87,7 @@ def symeval(val, symbols): """ first_replacement = {pystr_to_symbolic(k): pystr_to_symbolic('__REPLSYM_' + k) for k in symbols.keys()} second_replacement = {pystr_to_symbolic('__REPLSYM_' + k): v for k, v in symbols.items()} - return sp.simplify(val.subs(first_replacement).subs(second_replacement)) + return simplify(val.subs(first_replacement).subs(second_replacement)) def evaluate_symbols(base, new): @@ -137,31 +105,71 @@ def update_mapping(mapping, e): mapping.update(update) -def update_map_iterators(map, mapping): - # update the map params and return False - # if all iterations exhausted, return True - # always increase the last one. If it is exhausted, increase the next one and so forth +def assignment_misses(edge, mapping, stack, clt, C, symbols, array_names): + # regex pattern to detect buffer name and index if applicable + pattern = re.compile( + r""" + ^\s* + (?P[a-zA-Z_]\w*) # variable name + (?:\[ + (?P[^\[\]]+) # anything inside brackets (no nested []) + \])? + \s*$ +""", re.VERBOSE) + + misses = 0 + for lhs, rhs in edge.data.assignments.items(): + m_lhs = pattern.match(lhs) + m_rhs = pattern.match(rhs) + try: + lhs_name = m_lhs.group("name") + lhs_index = m_lhs.group("index") + if lhs_index and not lhs_index.isdigit(): + lhs_index = pystr_to_symbolic(m_lhs.group("index")) + elif lhs_index and lhs_index.isdigit(): + lhs_index = sp.Expr(int(lhs_index)) + + rhs_name = m_rhs.group("name") + rhs_index = m_rhs.group("index") + if rhs_index and not rhs_index.isdigit(): + rhs_index = pystr_to_symbolic(m_rhs.group("index")) + elif rhs_index and rhs_index.isdigit(): + lhs_index = sp.Expr(int(rhs_index)) + + if lhs_name in clt.array_info or (lhs_name in array_names and array_names[lhs_name] in clt.array_info): + line_id = clt.cache_line_id(lhs_name if lhs_name not in array_names else array_names[lhs_name], + ([lhs_index.subs(mapping)] if isinstance(lhs_index, sp.Expr) else []), + mapping) + line_id = int(line_id.subs(symbols).subs(mapping) if isinstance(line_id, sp.Expr) else line_id) + dist = stack.touch(line_id) + misses += 1 if dist >= C or dist == -1 else 0 + + if rhs_name in clt.array_info or (rhs_name in array_names and array_names[rhs_name] in clt.array_info): + line_id = clt.cache_line_id(rhs_name if rhs_name not in array_names else array_names[rhs_name], + ([rhs_index.subs(mapping)] if isinstance(rhs_index, sp.Expr) else []), + mapping) + line_id = int(line_id.subs(symbols).subs(mapping) if isinstance(line_id, sp.Expr) else line_id) + dist = stack.touch(line_id) + misses += 1 if dist >= C or dist == -1 else 0 + except Exception as e: + warnings.warn('Skipping a cache-miss contribution from an unparsable edge assignment: %s' % e) + return misses + + +def update_map_iterators(map, mapping, symbols): + # Advance the map iteration variables (innermost first); return True when all are exhausted. + # Increment the last param; if it overflows its range, reset it and carry to the next one. map_exhausted = True - for p, range in zip(map.params[::-1], map.range[::-1]): # reversed order - curr_value = mapping[p] - if not isinstance(range[1], SymExpr): - if curr_value.subs(mapping) + range[2].subs(mapping) <= range[1].subs(mapping): - # update this value and then we are done - mapping[p] = curr_value.subs(mapping) + range[2].subs(mapping) - map_exhausted = False - break - else: - # set current param to start again and continue - mapping[p] = range[0].subs(mapping) - else: - if curr_value.subs(mapping) + range[2].subs(mapping) <= range[1].expr.subs(mapping): - # update this value and we done - mapping[p] = curr_value.subs(mapping) + range[2].subs(mapping) - map_exhausted = False - break - else: - # set current param to start again and continue - mapping[p] = range[0].subs(mapping) + for p, rng in zip(map.params[::-1], map.range[::-1]): # reversed order + lo, hi, step = rng[0], (rng[1].expr if isinstance(rng[1], SymExpr) else rng[1]), rng[2] + curr = mapping[p].subs(symbols).subs(mapping) + step = step.subs(symbols).subs(mapping) + if curr + step <= hi.subs(symbols).subs(mapping): + mapping[p] = curr + step + map_exhausted = False + break + # this dimension is exhausted: reset it to its start and carry to the next + mapping[p] = lo.subs(symbols).subs(mapping) return map_exhausted @@ -174,25 +182,42 @@ def map_op_in(state: SDFGState, op_in_map: Dict[str, sp.Expr], entry, mapping, s map_misses = 0 while True: # do analysis of map contents - map_misses += scope_op_in(state, op_in_map, mapping, stack, clt, C, symbols, array_names, decided_branches, - ask_user, entry) + map_misses += scope_misses(state, op_in_map, mapping, stack, clt, C, symbols, array_names, decided_branches, + ask_user, entry) - if update_map_iterators(entry.map, mapping): + if update_map_iterators(entry.map, mapping, symbols): break return map_misses -def scope_op_in(state: SDFGState, - op_in_map: Dict[str, sp.Expr], - mapping, - stack: AccessStack, - clt: CacheLineTracker, - C, - symbols, - array_names, - decided_branches, - ask_user, - entry=None): +def _edge_miss(edge, clt: CacheLineTracker, array_names, mapping, symbols, stack: AccessStack, C) -> int: + """ + Account a single memlet access against the cache. + + :return: 1 if accessing the edge's data is a cache miss (or a first-ever touch), else 0. Edges + whose data is not a tracked global-memory array contribute nothing. + """ + data = edge.data.data + if data not in clt.array_info and not (data in array_names and array_names[data] in clt.array_info): + return 0 + line_id = clt.cache_line_id(data if data not in array_names else array_names[data], + [x[0].subs(mapping) for x in edge.data.subset.ranges], mapping) + line_id = int(line_id.subs(symbols).subs(mapping)) + dist = stack.touch(line_id) + return 1 if dist >= C or dist == -1 else 0 + + +def scope_misses(state: SDFGState, + op_in_map: Dict[str, sp.Expr], + mapping, + stack: AccessStack, + clt: CacheLineTracker, + C, + symbols, + array_names, + decided_branches, + ask_user, + entry=None): """ Computes the operational intensity of a single scope (scope is either an SDFG state or a map scope). @@ -222,19 +247,18 @@ def scope_op_in(state: SDFGState, update_map(op_in_map, get_uuid(node, state), map_misses) scope_misses += map_misses + elif isinstance(node, nd.AccessNode): + # A copy between two access nodes moves data without a tasklet, so the tasklet case below + # does not see it, yet it still touches memory. Only element-wise copies are accounted: + # _edge_miss models a single cache line touch, which says nothing about a bulk copy. + for e in state.out_edges(node): + if isinstance(e.dst, nd.AccessNode) and not e.data.is_empty() and e.data.subset.num_elements() == 1: + scope_misses += _edge_miss(e, clt, array_names, mapping, symbols, stack, C) elif isinstance(node, nd.Tasklet): tasklet_misses = 0 - # analyze the memory accesses of this tasklet and whether they hit in cache or not + # Account each tasklet memory access. for e in state.in_edges(node) + state.out_edges(node): - if e.data.data in clt.array_info or (e.data.data in array_names - and array_names[e.data.data] in clt.array_info): - line_id = clt.cache_line_id( - e.data.data if e.data.data not in array_names else array_names[e.data.data], - [x[0].subs(mapping) for x in e.data.subset.ranges], mapping) - - line_id = int(line_id.subs(mapping)) - dist = stack.touch(line_id) - tasklet_misses += 1 if dist >= C or dist == -1 else 0 + tasklet_misses += _edge_miss(e, clt, array_names, mapping, symbols, stack, C) scope_misses += tasklet_misses # a tasklet can get passed multiple times... we report the average misses in the end @@ -259,7 +283,7 @@ def scope_op_in(state: SDFGState, for e in state.out_edges(node): nested_array_names[e.src_conn] = e.data.data # Nested SDFGs are recursively analyzed first. - nsdfg_misses = sdfg_op_in(node.sdfg, op_in_map, mapping, stack, clt, C, nested_syms, nested_array_names, + nsdfg_misses = cfg_misses(node.sdfg, op_in_map, mapping, stack, clt, C, nested_syms, nested_array_names, decided_branches, ask_user) scope_misses += nsdfg_misses @@ -271,7 +295,7 @@ def scope_op_in(state: SDFGState, top_level_sdfg.add_symbol(f'{node.name}_misses', dtypes.int64) except FileExistsError: pass - lib_node_misses = sp.Symbol(f'{node.name}_misses', positive=True) + lib_node_misses = symbol(f'{node.name}_misses', positive=True) lib_node_misses = lib_node_misses.subs(mapping) scope_misses += lib_node_misses update_map(op_in_map, get_uuid(node, state), lib_node_misses) @@ -281,7 +305,133 @@ def scope_op_in(state: SDFGState, return scope_misses -def sdfg_op_in(sdfg: SDFG, +def cfr_misses(cfr: ControlFlowRegion, + op_in_map: Dict[str, Tuple[sp.Expr, sp.Expr]], + mapping, + stack: AccessStack, + clt: CacheLineTracker, + C, + symbols, + array_names, + decided_branches, + ask_user, + start=None): + region_misses = 0 + if isinstance(cfr, SDFGState): + region_misses = scope_misses(cfr, op_in_map, mapping, stack, clt, C, symbols, array_names, decided_branches, + ask_user, None) + + elif isinstance(cfr, LoopRegion): + loop_var = cfr.loop_variable + loop_condition = pystr_to_symbolic(cfr.loop_condition.as_string) + start = loop_analysis.get_init_assignment(cfr).subs(mapping) + step = pystr_to_symbolic(loop_analysis.get_loop_stride(cfr)) + mapping[loop_var] = start.subs(mapping) + region_misses = 0 + while (loop_condition.subs(mapping) == True): + iter_misses = cfg_misses(cfr, + op_in_map, + mapping, + stack, + clt, + C, + symbols, + array_names, + decided_branches, + ask_user, + start=cfr.start_block, + end=None) + mapping[loop_var] = mapping[loop_var] + step + region_misses += iter_misses + elif isinstance(cfr, ConditionalBlock): + true_branches = [] + possible_branches = [] + else_branch = None + + for cond, branch in cfr.branches: + if cond is None: + else_branch = branch + continue + + sym_cond = pystr_to_symbolic(cond.as_string) + res = sym_cond.subs(mapping) + + if res == True: + true_branches.append(branch) + elif res == False: + continue + else: + possible_branches.append(branch) + + ### if the branch is not decided by a true condition we + # 1- ask the userif he hasn't decided yet + # 2- take the one we took last time if he has decided + # 3- take the worst case if he opted not to decide + possibilities = true_branches + possible_branches + [else_branch] + if not true_branches and len(possible_branches) > 0 and ask_user and ( + cfr not in decided_branches or decided_branches[cfr] not in possibilities): + if len(possibilities) > 1: + print(f'\n\nWhich branch to take at {cfr.name}') + for i in range(len(possibilities)): + print(f'({i}) for branch {possibilities[i] if possibilities[i] else "else_branch"}') + chosen = int(input('Choose an option from above: ')) + # if the user chooses one, we check only that branch + branches = [possibilities[chosen]] + if possibilities[chosen]: + # only store the decided branch if it is not the implicit else branch + decided_branches[cfr] = possibilities[chosen] + else: + branches = possibilities + elif true_branches: + # if we have true branches we take the first one + branches = [true_branches[0]] + elif cfr in decided_branches and decided_branches[cfr] in possibilities: + # reuse the branch the user decided on a previous visit (e.g. a conditional inside a loop) + branches = [decided_branches[cfr]] + else: + # else we check all possibilities and take the max + branches = possibilities + + max_branch_misses = 0 + mapping_after_cond, stack_after_cond, decided_branches_after_cond = mapping, stack, decided_branches + for branch in branches: + if not branch: + # the implicit else branch has no misses + continue + # copy all data that must not be shared between branches + mapping_copy = deepcopy(mapping) + stack_copy = deepcopy(stack) + symbols_copy = deepcopy(symbols) + decided_branches_copy = deepcopy(decided_branches) + branch_misses = cfg_misses(branch, op_in_map, mapping_copy, stack_copy, clt, C, symbols_copy, array_names, + decided_branches_copy, ask_user, branch.start_block, None) + + if branch_misses > max_branch_misses: + max_branch_misses = branch_misses + mapping_after_cond, stack_after_cond, decided_branches_after_cond = mapping_copy, stack_copy, decided_branches_copy + + mapping.update(mapping_after_cond) + stack.replace_self(stack_after_cond) + decided_branches.update(decided_branches_after_cond) + region_misses = max_branch_misses + elif isinstance(cfr, FunctionCallRegion): + region_misses = cfg_misses(cfr, + op_in_map, + mapping, + stack, + clt, + C, + symbols, + array_names, + decided_branches, + ask_user, + start=cfr.start_block, + end=None) + + return region_misses + + +def cfg_misses(cfg: ControlFlowRegion, op_in_map: Dict[str, Tuple[sp.Expr, sp.Expr]], mapping, stack: AccessStack, @@ -311,112 +461,50 @@ def sdfg_op_in(sdfg: SDFG, :param end: The end state of the SDFG traversal. If None, the whole SDFG is traversed. """ - if start is None: + if isinstance(cfg, SDFG) and start is None: # add this SDFG's arrays to the cache line tracker - for name, arr in sdfg.arrays.items(): + for name, arr in cfg.arrays.items(): if isinstance(arr, Array): if name in array_names: name = array_names[name] clt.add_array(name, arr, mapping) # start traversal at SDFG's start state - curr_state = sdfg.start_state + curr_state = cfg.start_block else: curr_state = start total_misses = 0 - # traverse this SDFG's states + # traverse this SDFG's ControlFlowRegions while True: - total_misses += scope_op_in(curr_state, op_in_map, mapping, stack, clt, C, symbols, array_names, - decided_branches, ask_user) - if len(sdfg.out_edges(curr_state)) == 0: - # we reached an end state --> stop - break - else: - # take first edge with True condition - found = False - for e in sdfg.out_edges(curr_state): - if e.data.is_unconditional() or e.data.condition_sympy().subs(mapping) == True: - # save e's assignments in mapping and update curr_state - # replace values first with mapping, then update mapping - try: - update_mapping(mapping, e) - except: - print('\nWARNING: Uncommon assignment detected on InterstateEdge (e.g. bitwise operators).' - 'Analysis may give wrong results.') - print(e.data.assignments, 'was the edge\'s assignments.') - curr_state = e.dst - found = True - break - if not found: - # We need to check if we are in an implicit end state (i.e. all outgoing edge conditions evaluate to False) - all_false = True - for e in sdfg.out_edges(curr_state): - if e.data.condition_sympy().subs(mapping) != False: - all_false = False - if all_false: - break + region_misses = cfr_misses(curr_state, op_in_map, mapping, stack, clt, C, symbols, array_names, + decided_branches, ask_user) - if curr_state in decided_branches: - # if the user already decided this branch in a previous iteration, take the same branch again. - e = decided_branches[curr_state] - - update_mapping(mapping, e) - curr_state = e.dst - else: - # we cannot determine which branch to take --> check if both contain work - merge_state = find_merge_state(sdfg, curr_state) - next_edge_candidates = [] - for e in sdfg.out_edges(curr_state): - states = find_states_between(sdfg, e.dst, merge_state) - curr_work = mem_accesses_on_path(states) - if sp.sympify(curr_work).subs(mapping) > 0: - next_edge_candidates.append(e) - - if len(next_edge_candidates) == 1: - e = next_edge_candidates[0] - update_mapping(mapping, e) - decided_branches[curr_state] = e - curr_state = e.dst - else: - if ask_user: - edges = sdfg.out_edges(curr_state) - print(f'\n\nWhich branch to take at {curr_state.name}') - for i in range(len(edges)): - print(f'({i}) for edge to state {edges[i].dst.name}') - print(edges[i].dst._read_and_write_sets()) - print('merge state is named ', merge_state) - chosen = int(input('Choose an option from above: ')) - e = edges[chosen] - update_mapping(mapping, e) - decided_branches[curr_state] = e - curr_state = e.dst - print(2 * '\n') - else: - final_e = next_edge_candidates.pop() - for e in next_edge_candidates: - - # copy the state of the analysis - curr_mapping = dict(mapping) - update_mapping(curr_mapping, e) - curr_stack = stack.copy() - curr_clt = clt.copy() - curr_symbols = dict(symbols) - curr_array_names = dict(array_names) - - curr_state = e.dst - # walk down this branch until merge_state - sdfg_op_in(sdfg, op_in_map, curr_mapping, curr_stack, curr_clt, C, curr_symbols, - curr_array_names, decided_branches, ask_user, curr_state, merge_state) - - update_mapping(mapping, final_e) - curr_state = final_e.dst + total_misses += region_misses + out_edges = cfg.out_edges(curr_state) + if len(out_edges) == 0: + # reached an end state --> stop + break + # Structured control flow: a block has a single successor (branches are ConditionalBlocks, + # loops are LoopRegions). A statically-false condition on that edge is an implicit end. + edge = out_edges[0] + if not edge.data.is_unconditional() and edge.data.condition_sympy().subs(mapping) == False: + break + # Save the edge's assignments into the mapping and advance. + try: + total_misses += assignment_misses(edge, mapping, stack, clt, C, symbols, array_names) + update_mapping(mapping, edge) + except Exception: + warnings.warn('Uncommon assignment on an interstate edge (e.g. bitwise operators); ' + 'analysis may give wrong results. Assignments: %s' % edge.data.assignments) + curr_state = edge.dst if curr_state == end: break if end is None: # only update if we were actually analyzing a whole sdfg (not just start to end state) - update_map(op_in_map, get_uuid(sdfg), total_misses, average=False) + update_map(op_in_map, get_uuid(cfg), total_misses, average=False) + return total_misses @@ -427,11 +515,14 @@ def analyze_sdfg_op_in(sdfg: SDFG, assumptions, generate_plots=False, stringify=False, - test_set_size=3, + test_set_size=1, ask_user=False): """ Computes the operational intensity of the input SDFG. + :note: Only structured control flow is supported (loops as ``LoopRegion``, branches as + ``ConditionalBlock``, no ``break`` / ``continue`` / ``return``). An SDFG with unstructured + control flow is not analyzed: the analysis warns and returns a zero result. :param sdfg: The SDFG to analyze. :param op_in_map: Dictionary storing the resulting operational intensity for each SDFG element. :param C: Cache size in bytes. @@ -450,6 +541,16 @@ def analyze_sdfg_op_in(sdfg: SDFG, C = C // L sdfg = deepcopy(sdfg) + + # The analysis only models structured control flow. If the SDFG has a legacy loop or + # unstructured branching, bail out with a zero result rather than producing a wrong one. + if has_unstructured_control_flow(sdfg): + warnings.warn('Operational-intensity analysis supports only structured control flow ' + '(LoopRegion / ConditionalBlock); the SDFG contains a legacy loop or ' + 'unstructured branch, so no result is produced.') + op_in_map[get_uuid(sdfg)] = 0 + return + # apply SSA pass pipeline = FixedPointPipeline([StrictSymbolSSA()]) pipeline.apply_pass(sdfg, {}) @@ -463,9 +564,10 @@ def analyze_sdfg_op_in(sdfg: SDFG, elif isinstance(assumptions[sym], str): range_symbol[sym] = SymbolRange(int(x) for x in assumptions[sym].split(',')) del assumptions[sym] - work_map = {} + assumptions_list = [f'{x}=={y}' for x, y in assumptions.items()] + analyze_sdfg(sdfg, work_map, get_tasklet_work, assumptions_list) if len(undefined_symbols) > 0: @@ -480,12 +582,17 @@ def analyze_sdfg_op_in(sdfg: SDFG, # all symbols are concretized --> run normal op_in analysis with concretized symbols sdfg.specialize(assumptions) mapping = {} + # add the static symbols to the map to allow for better analysis + static_symbols = get_static_symbols(sdfg) + mapping.update(static_symbols) + mapping.update(assumptions) + mapping = {k: subs_till_fixed_point(v, mapping) for k, v in mapping.items()} stack = AccessStack(C) clt = CacheLineTracker(L) - sdfg_op_in(sdfg, op_in_map, mapping, stack, clt, C, {}, {}, {}, ask_user) + cfg_misses(sdfg, op_in_map, mapping, stack, clt, C, {}, {}, {}, ask_user) # compute bytes for k, v in op_in_map.items(): op_in_map[k] = v[0] / v[1] * L @@ -500,11 +607,12 @@ def analyze_sdfg_op_in(sdfg: SDFG, while True: new_val = False for sym, r in range_symbol.items(): + val = r.next() if val > -1: new_val = True assumptions[sym] = val - elif t < 3: + elif t < test_set_size: # now we sample test set t += 1 assumptions[sym] = r.max_value() + t * 3 @@ -512,12 +620,19 @@ def analyze_sdfg_op_in(sdfg: SDFG, if not new_val: break + r_sdfg = deepcopy(sdfg) + curr_op_in_map = {} mapping = {} + # add the static symbols to the map to allow for better analysis + static_symbols = get_static_symbols(r_sdfg) + mapping.update(static_symbols) mapping.update(assumptions) + mapping = {k: subs_till_fixed_point(v, mapping) for k, v in mapping.items()} + stack = AccessStack(C) clt = CacheLineTracker(L) - sdfg_op_in(sdfg, curr_op_in_map, mapping, stack, clt, C, {}, {}, {}, ask_user) + cfg_misses(r_sdfg, curr_op_in_map, mapping, stack, clt, C, {}, {}, {}, ask_user) # compute average cache misses for k, v in curr_op_in_map.items(): @@ -526,7 +641,7 @@ def analyze_sdfg_op_in(sdfg: SDFG, # save cache misses curr_cache_misses = dict(curr_op_in_map) - work_measurements.append(work_map[get_uuid(sdfg)][0].subs(assumptions)) + work_measurements.append(work_map[get_uuid(sdfg)].subs(assumptions)) # put curr values in cache_miss_measurements for k, v in curr_cache_misses.items(): if k in cache_miss_measurements: @@ -540,18 +655,16 @@ def analyze_sdfg_op_in(sdfg: SDFG, sympy_fs = {} for k, v in cache_miss_measurements.items(): + final_f, sympy_f, r_s = fit_curve(x_values[:-test_set_size], v[:-test_set_size], symbol_name) - op_in_map[k] = sp.simplify(sympy_f * L) + op_in_map[k] = simplify(sympy_f * L) sympy_fs[k] = sympy_f if k == get_uuid(sdfg): # compute MAPE on total SDFG mape = compute_mape(final_f, x_values[-test_set_size:], v[-test_set_size:], test_set_size) if mape > 0.2: - print('High MAPE detected:', mape) - print('It is suggested to generate plots and analyze those.') - print('R^2 is:', r_s) - print('A hight R^2 (i.e. close to 1) suggests that we are fitting the test data well.') - print('This combined with high MAPE tells us that our test data does not generalize.') + warnings.warn('High MAPE (%s) with R^2 = %s: the fit matches the test data but ' + 'may not generalize; generating plots is suggested.' % (mape, r_s)) calculate_op_in(op_in_map, work_map, not generate_plots) if generate_plots: @@ -562,6 +675,7 @@ def analyze_sdfg_op_in(sdfg: SDFG, if stringify: for k, v in op_in_map.items(): op_in_map[k] = str(v) + return op_in_map[get_uuid(sdfg)] ################################################################################ @@ -599,7 +713,6 @@ def main() -> None: assumptions[a] = int(b) else: assumptions[a] = b - print(assumptions) analyze_sdfg_op_in(sdfg, op_in_map, int(args.C), int(args.L), assumptions) result_whole_sdfg = op_in_map[get_uuid(sdfg)] diff --git a/dace/sdfg/performance_evaluation/work_depth.py b/dace/sdfg/performance_evaluation/work_depth.py index 35b60701c7..b59ad450c7 100644 --- a/dace/sdfg/performance_evaluation/work_depth.py +++ b/dace/sdfg/performance_evaluation/work_depth.py @@ -6,23 +6,27 @@ from collections import deque from dace.sdfg import nodes as nd, propagation, InterstateEdge from dace import SDFG, SDFGState, dtypes -from dace.subsets import Range -from typing import List, Tuple, Dict, Callable, Sequence, Union +from typing import List, Optional, Set, Tuple, Dict, Callable, Sequence, Union import os import sympy as sp from copy import deepcopy -from dace.libraries.blas import MatMul +from dace.libraries.blas import MatMul, Dot, Gemm, Gemv from dace.libraries.standard import Reduce -from dace.libraries.linalg import Transpose -from dace.symbolic import pystr_to_symbolic +from dace.libraries.linalg import Cholesky, Inv, Solve, Transpose +from dace.symbolic import pystr_to_symbolic, free_symbols_and_functions, symbol, int_floor, simplify import ast import astunparse import warnings -from dace.sdfg.performance_evaluation.helpers import LoopExtractionError, get_uuid, find_loop_guards_tails_exits +from dace.sdfg.performance_evaluation.helpers import get_uuid, get_static_symbols, has_unstructured_control_flow from dace.sdfg.performance_evaluation.assumptions import parse_assumptions from dace.transformation.passes.symbol_ssa import StrictSymbolSSA from dace.transformation.pass_pipeline import FixedPointPipeline +from dace.transformation.passes.analysis import loop_analysis + +from dace.sdfg.state import AbstractControlFlowRegion, ControlFlowRegion, LoopRegion, ConditionalBlock + +math_funcs = set() def get_array_size_symbols(sdfg): @@ -52,10 +56,17 @@ def symeval(val, symbols): """ first_replacement = {pystr_to_symbolic(k): pystr_to_symbolic('__REPLSYM_' + k) for k in symbols.keys()} second_replacement = {pystr_to_symbolic('__REPLSYM_' + k): v for k, v in symbols.items()} - return sp.simplify(val.subs(first_replacement).subs(second_replacement)) + return simplify(val.subs(first_replacement).subs(second_replacement)) def evaluate_symbols(base, new): + """Takes a base symbol mapping and a new one and adapts the new one to match the base one for symbols contained in it + + :param base: The base mapping + :param new: The mapping that gets adjusted + :return result: A new mapping that contains all mappings from new, but adjusted to transitively match to the mapping of base + """ + result = {} for k, v in new.items(): result[k] = symeval(pystr_to_symbolic(v), base) @@ -63,9 +74,9 @@ def evaluate_symbols(base, new): def count_work_matmul(node, symbols, state): - A_memlet = next(e for e in state.in_edges(node) if e.dst_conn == '_a') - B_memlet = next(e for e in state.in_edges(node) if e.dst_conn == '_b') - C_memlet = next(e for e in state.out_edges(node) if e.src_conn == '_c') + """Work of a matrix-multiply library node: 2*M*N*K flops, times the batch size if present.""" + A_memlet = next(state.in_edges_by_connector(node, '_a')) + C_memlet = next(state.out_edges_by_connector(node, '_c')) result = 2 # Multiply, add # Batch if len(C_memlet.data.subset) == 3: @@ -76,17 +87,18 @@ def count_work_matmul(node, symbols, state): result *= symeval(C_memlet.data.subset.size()[-1], symbols) # K result *= symeval(A_memlet.data.subset.size()[-1], symbols) - return sp.sympify(result) + return pystr_to_symbolic(result) def count_depth_matmul(node, symbols, state): - # optimal depth of a matrix multiplication is O(log(size of shared dimension)): - A_memlet = next(e for e in state.in_edges(node) if e.dst_conn == '_a') + """Depth of a matrix multiply: O(log K) over the shared (contracted) dimension K.""" + A_memlet = next(state.in_edges_by_connector(node, '_a')) size_shared_dimension = symeval(A_memlet.data.subset.size()[-1], symbols) - return sp.log(size_shared_dimension) + return sp.Max(1, sp.log(sp.Max(1, size_shared_dimension), 2)) def count_work_reduce(node, symbols, state): + """Work of a reduction library node: the WCR's arithmetic cost times the number of reduced elements.""" result = 0 if node.wcr is not None: result += count_arithmetic_ops_code(node.wcr) @@ -98,45 +110,276 @@ def count_work_reduce(node, symbols, state): result *= in_memlet.data.volume else: result = 0 - return sp.sympify(result) + return pystr_to_symbolic(result) def count_depth_reduce(node, symbols, state): - # optimal depth of reduction is log of the work - return sp.log(count_work_reduce(node, symbols, state)) + """Depth of a reduction: O(log(work)) for a balanced reduction tree.""" + return sp.log(sp.Max(1, count_work_reduce(node, symbols, state)), 2) + + +def count_work_dot(node, symbols, state): + """Work of a dot-product library node: 2*N - 1 flops (N multiplies and N-1 additions).""" + X_memlet = next(state.in_edges_by_connector(node, '_x')) + result = 2 * symeval(X_memlet.data.subset.size()[-1], symbols) - 1 + return pystr_to_symbolic(result) + + +def count_depth_dot(node, symbols, state): + """Depth of a dot product: one multiply layer plus O(log N) for the addition tree.""" + X_memlet = next(state.in_edges_by_connector(node, '_x')) + result = 1 + sp.log(sp.Max(1, symeval(X_memlet.data.subset.size()[-1], symbols)), 2) + return pystr_to_symbolic(result) + + +def count_work_cholesky(node, symbols, state): + """Work of a Cholesky factorization library node: ~N**3/3 flops for an N x N matrix.""" + A_memlet = next(state.in_edges_by_connector(node, '_a')) + N = symeval(A_memlet.data.subset.size()[-1], symbols) + return pystr_to_symbolic(N**3 / 3) + + +def count_depth_cholesky(node, symbols, state): + """Depth of a Cholesky factorization: N sequential elimination steps.""" + A_memlet = next(state.in_edges_by_connector(node, '_a')) + N = symeval(A_memlet.data.subset.size()[-1], symbols) + return sp.Max(1, N) + + +def count_work_inv(node, symbols, state): + """Work of a matrix-inverse library node: ~2*N**3 flops (LU ~2N**3/3 plus inversion ~4N**3/3).""" + A_memlet = next(state.in_edges_by_connector(node, '_ain')) + N = symeval(A_memlet.data.subset.size()[-1], symbols) + return pystr_to_symbolic(2 * N**3) + + +def count_depth_inv(node, symbols, state): + """Depth of a matrix inverse: N sequential elimination steps.""" + A_memlet = next(state.in_edges_by_connector(node, '_ain')) + N = symeval(A_memlet.data.subset.size()[-1], symbols) + return sp.Max(1, N) + + +def count_work_solve(node, symbols, state): + """Work of a linear-solve library node: LU (~2*N**3/3) plus forward/back substitution + (~2*N**2 per right-hand side).""" + A_memlet = next(state.in_edges_by_connector(node, '_ain')) + B_memlet = next(state.in_edges_by_connector(node, '_bin')) + N = symeval(A_memlet.data.subset.size()[-1], symbols) + b_size = B_memlet.data.subset.size() + rhs = symeval(b_size[-1], symbols) if len(b_size) >= 2 else 1 + return pystr_to_symbolic(2 * N**3 / 3 + 2 * N**2 * rhs) + + +def count_depth_solve(node, symbols, state): + """Depth of a linear solve: N sequential elimination steps.""" + A_memlet = next(state.in_edges_by_connector(node, '_ain')) + N = symeval(A_memlet.data.subset.size()[-1], symbols) + return sp.Max(1, N) + + +def count_work_gemm(node, symbols, state): + """ + Count work for GEMM operation: C = alpha * A @ B + beta * C + Work includes: + - Matrix multiplication: 2*M*N*K (multiply + add per element) + - Alpha scaling: M*N (if alpha != 1) + - Beta scaling + addition: 2*M*N (if beta != 0) + """ + A_memlet = next(state.in_edges_by_connector(node, '_a')) + C_memlet = next(state.out_edges_by_connector(node, '_c')) + + # Get dimensions + # Handle batch dimension if present + if len(C_memlet.data.subset) == 3: + batch = symeval(C_memlet.data.subset.size()[0], symbols) + M = symeval(C_memlet.data.subset.size()[1], symbols) + N = symeval(C_memlet.data.subset.size()[2], symbols) + else: + batch = 1 + M = symeval(C_memlet.data.subset.size()[-2], symbols) if len(C_memlet.data.subset.size()) >= 2 else 1 + N = symeval(C_memlet.data.subset.size()[-1], symbols) + + K = symeval(A_memlet.data.subset.size()[-1], symbols) + + # Core matrix multiplication: 2*M*N*K (multiply + add) + result = 2 * batch * M * N * K + + # Add work for alpha scaling if alpha != 1 + alpha = getattr(node, 'alpha', 1) + if alpha != 1: + result += batch * M * N # M*N multiplications by alpha + + # Add work for beta * C if beta != 0 + beta = getattr(node, 'beta', 0) + if beta != 0: + result += batch * M * N # M*N multiplications by beta + result += batch * M * N # M*N additions + + return pystr_to_symbolic(result) + + +def count_depth_gemm(node, symbols, state): + """ + Optimal depth for GEMM: log(K) for the reduction + constant for scaling/addition + """ + A_memlet = next(state.in_edges_by_connector(node, '_a')) + K = symeval(A_memlet.data.subset.size()[-1], symbols) + + # Depth is dominated by the reduction over K dimension + depth = sp.log(sp.Max(1, K), 2) + + # Add constant depth for alpha and beta operations + alpha = getattr(node, 'alpha', 1) + beta = getattr(node, 'beta', 0) + + if alpha != 1: + depth += 1 # One multiplication layer + if beta != 0: + depth += 2 # One multiplication + one addition layer + + return depth + + +def count_work_gemv(node, symbols, state): + """ + Count work for GEMV operation: y = alpha * A @ x + beta * y + Two variants: + - GEMV: y = alpha * A @ x + beta * y (A is MxN, x is N, y is M) + - GEMVT: y = alpha * A^T @ x + beta * y (A is MxN, x is M, y is N) + + Work includes: + - Matrix-vector multiplication: 2*M*N (multiply + add per element) + - Alpha scaling: M (if alpha != 1) + - Beta scaling + addition: 2*M (if beta != 0) + """ + A_memlet = next(state.in_edges_by_connector(node, '_A')) + + # Get dimensions from A matrix + A_shape = A_memlet.data.subset.size() + M = symeval(A_shape[-2], symbols) + N = symeval(A_shape[-1], symbols) + + # Check if transpose (GEMVT) + trans = getattr(node, 'transA', False) + + # Output size + output_size = N if trans else M + + # Core matrix-vector multiplication: 2*M*N (each output element needs N multiplies and N-1 adds) + result = 2 * M * N + + # Add work for alpha scaling if alpha != 1 + alpha = getattr(node, 'alpha', 1) + if alpha != 1: + result += output_size # output_size multiplications by alpha + + # Add work for beta * y if beta != 0 + beta = getattr(node, 'beta', 0) + if beta != 0: + result += output_size # output_size multiplications by beta + result += output_size # output_size additions + + return pystr_to_symbolic(result) + + +def count_depth_gemv(node, symbols, state): + """ + Optimal depth for GEMV: log(N) for the reduction + constant for scaling/addition + where N is the reduction dimension + """ + A_memlet = next(state.in_edges_by_connector(node, '_A')) + A_shape = A_memlet.data.subset.size() + M = symeval(A_shape[-2], symbols) + N = symeval(A_shape[-1], symbols) + + # Check if transpose + trans = getattr(node, 'transA', False) + + # Reduction dimension + reduction_dim = M if trans else N + + # Depth is dominated by the reduction + depth = sp.log(sp.Max(1, reduction_dim), 2) + + # Add constant depth for alpha and beta operations + alpha = getattr(node, 'alpha', 1) + beta = getattr(node, 'beta', 0) + + if alpha != 1: + depth += 1 # One multiplication layer + if beta != 0: + depth += 2 # One multiplication + one addition layer + + return sp.Max(1, depth) LIBNODES_TO_WORK = { MatMul: count_work_matmul, + Gemm: count_work_gemm, + Gemv: count_work_gemv, Transpose: lambda *args: 0, Reduce: count_work_reduce, + Dot: count_work_dot, + Cholesky: count_work_cholesky, + Inv: count_work_inv, + Solve: count_work_solve, } LIBNODES_TO_DEPTH = { MatMul: count_depth_matmul, + Gemm: count_depth_gemm, + Gemv: count_depth_gemv, Transpose: lambda *args: 0, Reduce: count_depth_reduce, + Dot: count_depth_dot, + Cholesky: count_depth_cholesky, + Inv: count_depth_inv, + Solve: count_depth_solve, } +# Type-cast calls (e.g. ``int``, ``float``, ``dace.float64``, ``dace.uint16``) perform no +# arithmetic and count as zero work. The names are derived from the available DaCe data types plus +# the Python/C builtins, so new dtypes need no maintenance here. +_TYPECAST_NAMES = ({'int', 'float', 'complex', 'bool', 'double'} + | {name + for name in dir(dtypes) if isinstance(getattr(dtypes, name), dtypes.typeclass)}) +_TYPECAST_NAMES |= {f'dace.{name}' for name in _TYPECAST_NAMES} + PYFUNC_TO_ARITHMETICS = { - 'float': 0, - 'dace.float64': 0, - 'dace.int64': 0, - 'dace.complex128': 0, + **{ + name: 0 + for name in _TYPECAST_NAMES + }, + # Transcendental intrinsics each count as one realised operation (np.* and math.* both lower to + # the bare C name in tasklet code); a user wanting hardware flop counts overrides these. 'math.exp': 1, 'exp': 1, + 'exp2': 1, 'math.tanh': 1, 'sin': 1, 'cos': 1, + 'tan': 1, + 'asin': 1, + 'acos': 1, + 'atan': 1, + 'atan2': 1, + 'sinh': 1, + 'cosh': 1, 'tanh': 1, + 'log': 1, + 'log2': 1, + 'log10': 1, 'math.sqrt': 1, 'sqrt': 1, - 'atan2': 1, + 'cbrt': 1, + 'int_floor': 1, # integer (floor) division + 'int_ceil': 1, # integer (ceil) division 'min': 0, 'max': 0, 'ceiling': 0, 'floor': 0, - 'abs': 0 + 'abs': 0, } @@ -159,9 +402,7 @@ def visit_UnaryOp(self, node): def visit_Call(self, node): fname = astunparse.unparse(node.func)[:-1] if fname not in PYFUNC_TO_ARITHMETICS: - print( - 'WARNING: Unrecognized python function "%s". If this is a type conversion, like "dace.float64", then this is fine.' - % fname) + warnings.warn('Unrecognized function "%s" in tasklet code; assuming it performs zero work.' % fname) return self.generic_visit(node) self.count += PYFUNC_TO_ARITHMETICS[fname] return self.generic_visit(node) @@ -224,9 +465,7 @@ def visit_UnaryOp(self, node): def visit_Call(self, node): fname = astunparse.unparse(node.func)[:-1] if fname not in PYFUNC_TO_ARITHMETICS: - print( - 'WARNING: Unrecognized python function "%s". If this is a type conversion, like "dace.float64", then this is fine.' - % fname) + warnings.warn('Unrecognized function "%s" in tasklet code; assuming it performs zero work.' % fname) # Still need to visit arguments to get their depth arg_depths = [self.visit(arg) for arg in node.args] return max(arg_depths) if arg_depths else 0 @@ -408,15 +647,15 @@ def tasklet_depth(tasklet_node: nd.Tasklet, state: SDFGState): def get_tasklet_work(node: nd.Tasklet, state: SDFGState): - return sp.sympify(tasklet_work(node, state)), sp.sympify(-1) + return pystr_to_symbolic(tasklet_work(node, state)), pystr_to_symbolic(-1) def get_tasklet_work_depth(node: nd.Tasklet, state: SDFGState): - return sp.sympify(tasklet_work(node, state)), sp.sympify(tasklet_depth(node, state)) + return pystr_to_symbolic(tasklet_work(node, state)), pystr_to_symbolic(tasklet_depth(node, state)) def get_tasklet_avg_par(node: nd.Tasklet, state: SDFGState): - return sp.sympify(tasklet_work(node, state)), sp.sympify(tasklet_depth(node, state)) + return pystr_to_symbolic(tasklet_work(node, state)), pystr_to_symbolic(tasklet_depth(node, state)) def update_value_map(old, new): @@ -433,24 +672,28 @@ def do_initial_subs(w, d, eq, subs1): """ Calls subs three times for the given (w)ork and (d)epth values. """ - result = sp.simplify(sp.sympify(w).subs(eq[0]).subs(eq[1]).subs(subs1)), sp.simplify( - sp.sympify(d).subs(eq[0]).subs(eq[1]).subs(subs1)) + result = simplify(pystr_to_symbolic(w).subs(eq[0]).subs(eq[1]).subs(subs1)), simplify( + pystr_to_symbolic(d).subs(eq[0]).subs(eq[1]).subs(subs1)) return result -def sdfg_work_depth(sdfg: SDFG, - w_d_map: Dict[str, Tuple[sp.Expr, sp.Expr]], - analyze_tasklet: Callable[[nd.Tasklet, SDFGState], Tuple[sp.Expr, sp.Expr]], - symbols: Dict[str, str], - equality_subs: Tuple[Dict[str, sp.Symbol], Dict[str, sp.Expr]], - subs1: Dict[str, sp.Expr], - detailed_analysis: bool = False) -> Tuple[sp.Expr, sp.Expr]: +def control_flow_region_work_depth( + cfr: ControlFlowRegion, + w_d_map: Dict[str, Tuple[sp.Expr | List[Tuple[sp.Expr, sp.Expr]], sp.Expr | List[Tuple[sp.Expr, sp.Expr]]]], + analyze_tasklet: Callable[[nd.Tasklet, SDFGState], Tuple[sp.Expr, sp.Expr]], + symbols: Dict[str, str], + equality_subs: Tuple[Dict[str, sp.Symbol], Dict[str, sp.Expr]], + subs1: Dict[str, sp.Expr], + detailed_analysis: bool = False, + data_symbols: Optional[Set[str]] = None +) -> Tuple[sp.Expr | List[Tuple[sp.Expr, sp.Expr]], sp.Expr | List[Tuple[sp.Expr, sp.Expr]]]: """ - Analyze the work and depth of a given SDFG. - First we determine the work and depth of each state. Then we break loops in the state machine, such that we get a DAG. - Lastly, we compute the path with most work and the path with the most depth in order to get the total work depth. + Analyze the work and depth of a given (structured) ControlFlowRegion. + First we determine the work and depth of each block (loops are ``LoopRegion`` and branches are + ``ConditionalBlock``, so the region itself is a DAG). Then we compute the path with the most work + and the path with the most depth to get the region's total work and depth. - :param sdfg: The SDFG to analyze. + :param cfr: The ControlFLowRegion to analyze. :param w_d_map: Dictionary which will save the result. :param analyze_tasklet: Function used to analyze tasklet nodes. :param symbols: A dictionary mapping local nested SDFG symbols to global symbols. @@ -459,100 +702,189 @@ def sdfg_work_depth(sdfg: SDFG, as computation time sky-rockets, since expression can became HUGE (depending on number of branches etc.). :param equality_subs: Substitution dict taking care of the equality assumptions. :param subs1: First substitution dict for greater/lesser assumptions. + :param data_symbols: The compute symbols of the owning SDFG (see :func:`compute_symbols`); computed once and + reused across the recursion. Interstate-edge arithmetic only counts as work for these. :return: A tuple containing the work and depth of the SDFG. """ + if data_symbols is None: + data_symbols = compute_symbols(cfr if isinstance(cfr, SDFG) else cfr.sdfg) + + # First determine the work and depth of each ControlFlowRegion individually. + # Keep track of the work and depth for each state in a dictionary + region_depths: Dict[AbstractControlFlowRegion, sp.Expr] = {} + region_works: Dict[AbstractControlFlowRegion, sp.Expr] = {} + for region in cfr.nodes(): + if isinstance(region, SDFGState): + #rename variable to make code more readable + state = region + state_work, state_depth = state_work_depth(state, w_d_map, analyze_tasklet, symbols, equality_subs, subs1, + detailed_analysis) + # Substitutions for state_work and state_depth already performed, but state.executions needs to be subs'd now. + state_work = simplify(state_work.subs(equality_subs[0]).subs(equality_subs[1]).subs(subs1)) + state_depth = simplify(state_depth.subs(equality_subs[0]).subs(equality_subs[1]).subs(subs1)) + + region_works[state], region_depths[state] = state_work, state_depth + w_d_map[get_uuid(state)] = (region_works[state], region_depths[state]) + elif isinstance(region, LoopRegion): + #rename variable to make code more readable + loop = region + fallback = False + + # We try to get a closed form solution for the work and depth of the loop. If this is not possible (e.g. because there are no static loop bounds), + # we fall back to just multiplying the work and depth of one loop iteration with the number of loop iterations. This can lead to incorrect results, + # since it does not take into account that work and depth of one loop iteration might be dependent on the loop variable. + # An unbounded loop may have no loop variable; ``None`` then forces the fallback below. + loop_var = symbol(loop.loop_variable) if loop.loop_variable else None + lower_bound = loop_analysis.get_init_assignment(loop) + upper_bound = loop_analysis.get_loop_end(loop) + step = pystr_to_symbolic(loop_analysis.get_loop_stride(loop)) + if any(v is None for v in (loop_var, lower_bound, upper_bound)): + warnings.warn('Loop without a static loop variable/bounds; falling back to its ' + 'execution count, which can affect the resulting expression.') + fallback = True + executions = loop.start_block.executions + executions = executions.subs(equality_subs[0]).subs(equality_subs[1]).subs(subs1) + + # Recursively get the work and depth of the loop body + loop_work, loop_depth = control_flow_region_work_depth(loop, w_d_map, analyze_tasklet, symbols, + equality_subs, subs1, detailed_analysis, + data_symbols) + + if not fallback: + # If static loop bounds are available, we can write the work and depth of the loop as a summation over the loop variable from the lower to the upper bound. + + # to ensure that the summation works properly, we need to make sure that the symbol that is used as loop varaible + # is the same as the ones used in the inner expression + for var in loop_work.free_symbols: + if var.name == loop_var.name: + loop_var = var + + #TEMPORARY FIX: with library nodes it can happen that we get two symbols (with the same name) + for var in loop_work.free_symbols: + if var.name == loop_var.name and not var == loop_var: + loop_work = loop_work.subs({var: loop_var}) + loop_depth = loop_depth.subs({var: loop_var}) + + # Accumulate the per-iteration work and depth over the loop range (shared with the + # map handler), so iteration-dependent work is summed rather than multiplied. + loop_work = accumulate_over_range(loop_work, loop_var, lower_bound, upper_bound, step, equality_subs, + subs1) + loop_depth = accumulate_over_range(loop_depth, loop_var, lower_bound, upper_bound, step, equality_subs, + subs1) + + # Do equality subs + loop_work = simplify(loop_work.subs(equality_subs[0]).subs(equality_subs[1]).subs(subs1)) + loop_depth = simplify(loop_depth.subs(equality_subs[0]).subs(equality_subs[1]).subs(subs1)) + else: + loop_work = simplify(loop_work.subs(equality_subs[0]).subs(equality_subs[1]).subs(subs1)) + loop_depth = simplify(loop_depth.subs(equality_subs[0]).subs(equality_subs[1]).subs(subs1)) - # First determine the work and depth of each state individually. - # Keep track of the work and depth for each state in a dictionary, where work and depth are multiplied by the number - # of times the state will be executed. - state_depths: Dict[SDFGState, sp.Expr] = {} - state_works: Dict[SDFGState, sp.Expr] = {} - for state in sdfg.nodes(): - state_work, state_depth = state_work_depth(state, w_d_map, analyze_tasklet, symbols, equality_subs, subs1, - detailed_analysis) + if executions != 0: + loop_work = loop_work * executions + loop_depth = loop_depth * executions + else: + exec_symbol = symbol(f'num_execs_{region.sdfg.cfg_id}_{region.sdfg.node_id(region)}', + nonnegative=True) + loop_work = loop_work * exec_symbol + loop_depth = loop_depth * exec_symbol + + region_works[loop], region_depths[loop] = loop_work, loop_depth + w_d_map[get_uuid(loop)] = (region_works[loop], region_depths[loop]) + + elif isinstance(region, ConditionalBlock): + branch_conditions = {} + branch_works = {} + branch_depths = {} + for (condition, branch) in region.branches: + branch_conditions[branch] = (pystr_to_symbolic(condition.as_string) + if condition is not None else pystr_to_symbolic(True)) + branch_works[branch], branch_depths[branch] = control_flow_region_work_depth( + branch, w_d_map, analyze_tasklet, symbols, equality_subs, subs1, detailed_analysis, data_symbols) + + if analyze_tasklet == get_tasklet_avg_par: + # For avg_par we want the branch minimising W/D (worst case parallelism). + # Build a Piecewise that selects work and depth together based on which + # branch has the smallest W/D ratio, so the pair stays consistent. + branches = list(branch_works.keys()) + + def avg_par_expr_val(w: sp.Expr, d: sp.Expr) -> sp.Expr: + """Return W/D, treating D=0 as infinite parallelism (never worst-case).""" + return sp.Piecewise((w / d, d > 0), (sp.oo, True)) + + # Start with the first branch as the running minimum + best_work = branch_works[branches[0]] + best_depth = branch_depths[branches[0]] + + for b in branches[1:]: + w_b = branch_works[b] + d_b = branch_depths[b] + ap_best = avg_par_expr_val(best_work, best_depth) + ap_b = avg_par_expr_val(w_b, d_b) + is_worse = simplify(ap_b < ap_best) # lower ratio = worse parallelism + best_work = sp.Piecewise((w_b, is_worse), (best_work, True)) + best_depth = sp.Piecewise((d_b, is_worse), (best_depth, True)) + + region_works[region] = best_work + region_depths[region] = best_depth + elif not detailed_analysis: + region_works[region] = sp.Max(*branch_works.values()) + region_depths[region] = sp.Max(*branch_depths.values()) + else: + work_condition = list(zip(branch_works.values(), branch_conditions.values())) + depth_condition = list(zip(branch_depths.values(), branch_conditions.values())) + region_works[region] = sp.Piecewise(*work_condition) + region_depths[region] = sp.Piecewise(*depth_condition) + w_d_map[get_uuid(region)] = (region_works[region], region_depths[region]) - # Substitutions for state_work and state_depth already performed, but state.executions needs to be subs'd now. - state_work = sp.simplify( - state_work.subs(equality_subs[0]).subs(equality_subs[1]).subs(subs1) * - state.executions.subs(equality_subs[0]).subs(equality_subs[1]).subs(subs1)) - state_depth = sp.simplify( - state_depth.subs(equality_subs[0]).subs(equality_subs[1]).subs(subs1) * - state.executions.subs(equality_subs[0]).subs(equality_subs[1]).subs(subs1)) + else: + function_work, function_depth = control_flow_region_work_depth(region, w_d_map, analyze_tasklet, symbols, + equality_subs, subs1, detailed_analysis, + data_symbols) + function_work = simplify(function_work.subs(equality_subs[0]).subs(equality_subs[1]).subs(subs1)) + function_depth = simplify(function_depth.subs(equality_subs[0]).subs(equality_subs[1]).subs(subs1)) - state_works[state], state_depths[state] = state_work, state_depth - w_d_map[get_uuid(state)] = (state_works[state], state_depths[state]) + region_works[region], region_depths[region] = function_work, function_depth + w_d_map[get_uuid(region)] = (region_works[region], region_depths[region]) edge_w_d_map: Dict[Tuple[str, str], Tuple[sp.Expr, sp.Expr]] = {} - for isedge in sdfg.edges(): - edge_work, edge_depth = sp.sympify(0), sp.sympify(0) + for isedge in cfr.edges(): + edge_work, edge_depth = pystr_to_symbolic(0), pystr_to_symbolic(0) if isedge.data.assignments: - for v in isedge.data.assignments.values(): + for sym, v in isedge.data.assignments.items(): edge_depth = sp.Max(edge_depth, count_depth_code(v)) - edge_work += count_arithmetic_ops_code(v) + # Only count arithmetic that computes a value consumed by computation; arithmetic + # assigned to addressing-only symbols (indices, bounds) is address computation, not work. + if sym in data_symbols: + edge_work += count_arithmetic_ops_code(v) edge_w_d_map[(get_uuid(isedge.src), get_uuid(isedge.dst))] = (edge_work, edge_depth) - # Prepare the SDFG for a depth analysis by breaking loops. This removes the edge between the last loop state and - # the guard, and instead places an edge between the last loop state and the exit state. - # This transforms the state machine into a DAG. Hence, we can find the "heaviest" and "deepest" paths in linear time. - # Additionally, construct a dummy exit state and connect every state that has no outgoing edges to it. - - # identify all loops in the SDFG - try: - nodes_oNodes_exits = find_loop_guards_tails_exits(sdfg._nx) - except LoopExtractionError: - # If loop detection fails, we cannot make proper propagation. - print('Analysis failed since not all loops got detected. It may help to use more structured loop constructs.' + - ' The analysis per state remains correct, but no SDFG-wide analysis can be performed.') - sdfg_result = (sp.oo, sp.oo) - w_d_map[get_uuid(sdfg)] = sdfg_result - - for k, (v_w, v_d) in w_d_map.items(): - # The symeval replaces nested SDFG symbols with their global counterparts. - v_w = symeval(v_w, symbols) - v_d = symeval(v_d, symbols) - w_d_map[k] = (v_w, v_d) - return sdfg_result - - # Now we need to go over each triple (node, oNode, exits). For each triple, we - # - remove edge (oNode, node), i.e. the backward edge - # - for all exits e, add edge (oNode, e). This edge may already exist - # - remove edge from node to exit (if present, i.e. while-do loop) - # - This ensures that every node with > 1 outgoing edge is a branch guard - # - useful for detailed anaylsis. - for node, oNode, exits in nodes_oNodes_exits: - sdfg.remove_edge(sdfg.edges_between(oNode, node)[0]) - for e in exits: - if len(sdfg.edges_between(oNode, e)) == 0: - # no edge there yet - sdfg.add_edge(oNode, e, InterstateEdge()) - if len(sdfg.edges_between(node, e)) > 0: - # edge present --> remove it - sdfg.remove_edge(sdfg.edges_between(node, e)[0]) - - # add a dummy exit to the SDFG, such that each path ends there. - dummy_exit = sdfg.add_state('dummy_exit') - for state in sdfg.nodes(): - if len(sdfg.out_edges(state)) == 0 and state is not dummy_exit: - sdfg.add_edge(state, dummy_exit, InterstateEdge()) + # Add a dummy exit so every path ends there. The analysis assumes structured control flow, so + # loops are LoopRegions (single nodes here) and this control-flow region is already a DAG; the + # BFS below can find the heaviest/deepest paths in linear time. + dummy_exit = cfr.add_state('dummy_exit') + for region in cfr.nodes(): + if len(cfr.out_edges(region)) == 0 and region is not dummy_exit: + cfr.add_edge(region, dummy_exit, InterstateEdge()) # These two dicts save the current length of the "heaviest", resp. "deepest", paths at each state. - work_map: Dict[SDFGState, sp.Expr] = {} - depth_map: Dict[SDFGState, sp.Expr] = {} + work_map: Dict[AbstractControlFlowRegion, sp.Expr] = {} + depth_map: Dict[AbstractControlFlowRegion, sp.Expr] = {} # Keeps track of assignments done on InterstateEdges. - state_value_map: Dict[SDFGState, Dict[sp.Symbol, sp.Symbol]] = {} + region_value_map: Dict[AbstractControlFlowRegion, Dict[sp.Symbol, sp.Symbol]] = {} # The dummy state has 0 work and depth. - state_depths[dummy_exit] = sp.sympify(0) - state_works[dummy_exit] = sp.sympify(0) + region_depths[dummy_exit] = pystr_to_symbolic(0) + region_works[dummy_exit] = pystr_to_symbolic(0) # Perform a BFS traversal of the state machine and calculate the maximum work / depth at each state. Only advance to # the next state in the BFS if all incoming edges have been visited, to ensure the maximum work / depth expressions # have been calculated. traversal_q = deque() - traversal_q.append((sdfg.start_state, sp.sympify(0), sp.sympify(0), None, [], [], {})) + traversal_q.append((cfr.start_block, pystr_to_symbolic(0), pystr_to_symbolic(0), None, [], [], {})) visited = set() - + c = 0 while traversal_q: - state, depth, work, ie, condition_stack, common_subexpr_stack, value_map = traversal_q.popleft() + c += 1 + region, depth, work, ie, condition_stack, common_subexpr_stack, value_map = traversal_q.popleft() if ie is not None: visited.add(ie) @@ -561,43 +893,43 @@ def sdfg_work_depth(sdfg: SDFG, work += edge_w_d_map[edge_uid][0] depth += edge_w_d_map[edge_uid][1] - if state in state_value_map: + if region in region_value_map: # update value map: - update_value_map(state_value_map[state], value_map) + update_value_map(region_value_map[region], value_map) else: - state_value_map[state] = value_map + region_value_map[region] = value_map - value_map = {pystr_to_symbolic(k): pystr_to_symbolic(v) for k, v in state_value_map[state].items()} - n_depth = sp.simplify((depth + state_depths[state]).subs(value_map)) - n_work = sp.simplify((work + state_works[state]).subs(value_map)) + value_map = {pystr_to_symbolic(k): pystr_to_symbolic(v) for k, v in region_value_map[region].items()} + n_depth = simplify((depth + region_depths[region]).subs(value_map)) + n_work = simplify((work + region_works[region]).subs(value_map)) # If we are analysing average parallelism, we don't search "heaviest" and "deepest" paths separately, but we want one # single path with the least average parallelsim (of all paths with more than 0 work). if analyze_tasklet == get_tasklet_avg_par: - if state in depth_map: # this means we have already visited this state before + if region in depth_map: # this means we have already visited this region before cse = common_subexpr_stack.pop() # if current path has 0 depth (--> 0 work as well), we don't do anything. if n_depth != 0: # check if we need to update the work and depth of the current state # we update if avg parallelism of new incoming path is less than current avg parallelism - if depth_map[state] == 0: + if depth_map[region] == 0: # old value was divided by zero --> we take new value anyway - depth_map[state] = cse[1] + n_depth - work_map[state] = cse[0] + n_work + depth_map[region] = cse[1] + n_depth + work_map[region] = cse[0] + n_work else: - old_avg_par = (cse[0] + work_map[state]) / (cse[1] + depth_map[state]) + old_avg_par = (cse[0] + work_map[region]) / (cse[1] + depth_map[region]) new_avg_par = (cse[0] + n_work) / (cse[1] + n_depth) # we take either old work/depth or new work/depth (or both if we cannot determine which one is greater) - depth_map[state] = cse[1] + sp.Piecewise((n_depth, sp.simplify(new_avg_par < old_avg_par)), - (depth_map[state], True)) - work_map[state] = cse[0] + sp.Piecewise((n_work, sp.simplify(new_avg_par < old_avg_par)), - (work_map[state], True)) + depth_map[region] = cse[1] + sp.Piecewise((n_depth, simplify(new_avg_par < old_avg_par)), + (depth_map[region], True)) + work_map[region] = cse[0] + sp.Piecewise((n_work, simplify(new_avg_par < old_avg_par)), + (work_map[region], True)) else: - depth_map[state] = n_depth - work_map[state] = n_work + depth_map[region] = n_depth + work_map[region] = n_work else: # search heaviest and deepest path separately - if state in depth_map: # and consequently also in work_map + if region in depth_map: # and consequently also in work_map # This cse value would appear in both arguments of the Max. Hence, for performance reasons, # we pull it out of the Max expression. # Example: We do cse + Max(a, b) instead of Max(cse + a, cse + b). @@ -607,18 +939,18 @@ def sdfg_work_depth(sdfg: SDFG, if detailed_analysis: # This MAX should be covered in the more detailed analysis cond = condition_stack.pop() - work_map[state] = cse[0] + sp.Piecewise((work_map[state], sp.Not(cond)), (n_work, cond)) - depth_map[state] = cse[1] + sp.Piecewise((depth_map[state], sp.Not(cond)), (n_depth, cond)) + work_map[region] = cse[0] + sp.Piecewise((work_map[region], sp.Not(cond)), (n_work, cond)) + depth_map[region] = cse[1] + sp.Piecewise((depth_map[region], sp.Not(cond)), (n_depth, cond)) else: - work_map[state] = cse[0] + sp.Max(work_map[state], n_work) - depth_map[state] = cse[1] + sp.Max(depth_map[state], n_depth) + work_map[region] = cse[0] + sp.Max(work_map[region], n_work) + depth_map[region] = cse[1] + sp.Max(depth_map[region], n_depth) else: - depth_map[state] = n_depth - work_map[state] = n_work + depth_map[region] = n_depth + work_map[region] = n_work - out_edges = sdfg.out_edges(state) + out_edges = cfr.out_edges(region) # only advance after all incoming edges were visited (meaning that current work depth values of state are final). - if any(iedge not in visited for iedge in sdfg.in_edges(state)): + if any(iedge not in visited for iedge in cfr.in_edges(region)): pass else: for oedge in out_edges: @@ -629,9 +961,9 @@ def sdfg_work_depth(sdfg: SDFG, new_cond_stack.append(oedge.data.condition_sympy()) # same for common_subexr_stack new_cse_stack = list(common_subexpr_stack) - new_cse_stack.append((work_map[state], depth_map[state])) + new_cse_stack.append((work_map[region], depth_map[region])) # same for value_map - new_value_map = dict(state_value_map[state]) + new_value_map = dict(region_value_map[region]) new_value_map.update({ pystr_to_symbolic(k): pystr_to_symbolic(v).subs(equality_subs[0]).subs(equality_subs[1]).subs(subs1) @@ -639,13 +971,12 @@ def sdfg_work_depth(sdfg: SDFG, }) traversal_q.append((oedge.dst, 0, 0, oedge, new_cond_stack, new_cse_stack, new_value_map)) else: - # value_map.update(oedge.data.assignments) value_map.update({ pystr_to_symbolic(k): pystr_to_symbolic(v).subs(equality_subs[0]).subs(equality_subs[1]).subs(subs1) for k, v in oedge.data.assignments.items() }) - traversal_q.append((oedge.dst, depth_map[state], work_map[state], oedge, condition_stack, + traversal_q.append((oedge.dst, depth_map[region], work_map[region], oedge, condition_stack, common_subexpr_stack, value_map)) try: @@ -654,18 +985,90 @@ def sdfg_work_depth(sdfg: SDFG, except KeyError: # If we get a KeyError above, this means that the traversal never reached the dummy_exit state. # This happens if the loops were not properly detected and broken. - raise LoopExtractionError( - 'Analysis failed, since not all loops got detected. It may help to use more structured loop constructs.') + raise RuntimeError("Analysis failed! The dummy exit state was never reached") - sdfg_result = (max_work, max_depth) - w_d_map[get_uuid(sdfg)] = sdfg_result + cfr_result = (max_work.simplify(), max_depth.simplify()) + w_d_map[get_uuid(cfr)] = cfr_result for k, (v_w, v_d) in w_d_map.items(): # The symeval replaces nested SDFG symbols with their global counterparts. v_w = symeval(v_w, symbols) v_d = symeval(v_d, symbols) w_d_map[k] = (v_w, v_d) - return sdfg_result + + return cfr_result + + +def compute_symbols(sdfg: SDFG) -> Set[str]: + """ + Return the names of symbols whose value is consumed by computation (as opposed to addressing). + + A symbol is a compute symbol if it is read inside a tasklet's code, or if it (transitively, + through interstate-edge assignments) feeds such a symbol. The complement -- symbols used only in + memlet subsets, map ranges and loop/branch conditions -- are addressing symbols. This lets the + work analysis attribute interstate-edge assignment arithmetic to computation rather than to + address calculation (e.g. the bitwise arithmetic on the loop-carried scalars of a CRC kernel is + compute, whereas an ``idx = j * N`` index helper is not). + + :param sdfg: The SDFG to inspect. + :return: The set of compute-symbol names. + """ + data_symbols: Set[str] = set() + for node, _ in sdfg.all_nodes_recursive(): + if isinstance(node, nd.Tasklet): + data_symbols |= {str(s) for s in node.free_symbols} + # Transitive closure: if a compute symbol is assigned an expression, the symbols feeding that + # expression are compute symbols too. + assignment_edges = [ + e for e, _ in sdfg.all_edges_recursive() if isinstance(e.data, InterstateEdge) and e.data.assignments + ] + changed = True + while changed: + changed = False + for edge in assignment_edges: + for lhs, rhs in edge.data.assignments.items(): + if lhs in data_symbols: + for sym in free_symbols_and_functions(rhs): + if sym not in data_symbols: + data_symbols.add(sym) + changed = True + return data_symbols + + +def accumulate_over_range(expr: sp.Expr, var: sp.Symbol, lower: sp.Expr, upper: sp.Expr, step: sp.Expr, + equality_subs: Tuple[Dict[str, sp.Symbol], + Dict[str, sp.Expr]], subs1: Dict[str, sp.Expr]) -> sp.Expr: + """ + Accumulate ``expr`` over one map/loop dimension ``var`` ranging over ``lower:upper:step`` (with + ``upper`` inclusive). Shared by the loop and map handlers so both accumulate identically. + + The summation is written as ``Sum(expr[var -> step*var + lower], (var, 0, (upper-lower)//step))``, + which both sums iteration-dependent work (e.g. the inner bound of a triangular nest) and reduces + to a multiplication when ``expr`` does not depend on ``var``. The iteration symbol is first + aligned with the assumption-substituted symbol that appears in ``expr`` (via ``subs1``). + + :param expr: The per-iteration work or depth expression. + :param var: The iteration variable. + :param lower: Inclusive lower bound of the iteration range. + :param upper: Inclusive upper bound of the iteration range. + :param step: Iteration stride. + :param equality_subs: Substitution dicts for the equality assumptions. + :param subs1: Substitution dict for the greater/lesser assumptions. + :return: The accumulated expression over the range. + """ + lower, upper, step = pystr_to_symbolic(lower), pystr_to_symbolic(upper), pystr_to_symbolic(step) + # Align the iteration symbol with the (assumption-substituted) symbol used inside ``expr``. + var = var.subs(subs1) + for sym in expr.free_symbols: + if sym.name == var.name and sym != var: + expr = expr.subs({sym: var}) + shifted_hi = int_floor(upper - lower, step).subs(equality_subs[0]).subs(equality_subs[1]).subs(subs1) + # Iterate from the lower bound unless the step is known-negative (a symbolic step, e.g. a tile + # size, is treated as forward; map steps are never negative). + lower = lower.subs(subs1) if step.is_negative is not True else upper.subs(subs1) + step = sp.Abs(step) + expr = expr.subs({var: step * var + lower}) + return sp.Sum(expr, (var, pystr_to_symbolic(0), shifted_hi)).doit() def scope_work_depth( @@ -704,13 +1107,13 @@ def scope_work_depth( # find the work and depth of each node # for maps and nested SDFG, we do it recursively - work = sp.sympify(0) - max_depth = sp.sympify(0) + work = pystr_to_symbolic(0) + max_depth = pystr_to_symbolic(0) scope_nodes = state.scope_children()[entry] scope_exit = None if entry is None else state.exit_node(entry) for node in scope_nodes: # add node to map - w_d_map[get_uuid(node, state)] = (sp.sympify(0), sp.sympify(0)) + w_d_map[get_uuid(node, state)] = (pystr_to_symbolic(0), pystr_to_symbolic(0)) if isinstance(node, nd.EntryNode): # If the scope contains an entry node, we need to recursively analyze the sub-scope of the entry node first. # The resulting work/depth are summarized into the entry node @@ -737,10 +1140,14 @@ def scope_work_depth( nested_syms.update(symbols) nested_syms.update(evaluate_symbols(symbols, node.symbol_mapping)) # Nested SDFGs are recursively analyzed first. - nsdfg_work, nsdfg_depth = sdfg_work_depth(node.sdfg, w_d_map, analyze_tasklet, nested_syms, equality_subs, - subs1, detailed_analysis) + nsdfg_work, nsdfg_depth = control_flow_region_work_depth(node.sdfg, w_d_map, analyze_tasklet, {}, + equality_subs, {}, detailed_analysis) + nsdfg_work, nsdfg_depth = nsdfg_work.subs(nested_syms), nsdfg_depth.subs( + nested_syms + ) # We cannot use assumptions for nested sdfg analysis. It interfers with the global assumptions. We thus substitute afterwards nsdfg_work, nsdfg_depth = do_initial_subs(nsdfg_work, nsdfg_depth, equality_subs, subs1) + # add up work for whole state, but also save work for this nested SDFG in w_d_map work += nsdfg_work w_d_map[get_uuid(node, state)] = (nsdfg_work, nsdfg_depth) @@ -759,33 +1166,38 @@ def scope_work_depth( # Such a library node was already encountered by the analysis. # Hence, we don't need to add anyting. pass - lib_node_work = sp.Symbol(f'{node.name}_work', positive=True) - lib_node_depth = sp.sympify(-1) + lib_node_work = symbol(f'{node.name}_work', positive=True) + lib_node_depth = pystr_to_symbolic(-1) if analyze_tasklet != get_tasklet_work: # we are analyzing depth try: lib_node_depth = LIBNODES_TO_DEPTH[type(node)](node, symbols, state) except KeyError: top_level_sdfg = state.parent - top_level_sdfg.add_symbol(f'{node.name}_depth', dtypes.int64) - lib_node_depth = sp.Symbol(f'{node.name}_depth', positive=True) + try: + top_level_sdfg.add_symbol(f'{node.name}_depth', dtypes.int64) + except FileExistsError: + pass + lib_node_depth = symbol(f'{node.name}_depth', positive=True) lib_node_work, lib_node_depth = do_initial_subs(lib_node_work, lib_node_depth, equality_subs, subs1) work += lib_node_work w_d_map[get_uuid(node, state)] = (lib_node_work, lib_node_depth) if entry is not None: - # If the scope being analyzed is a map, multiply the work by the number of iterations of the map. + # If the scope being analyzed is a map, accumulate the body work over its iteration domain. + # We accumulate per dimension (summing when the work depends on that map parameter, else + # multiplying), so that work depending on an enclosing iteration variable -- e.g. a + # triangular map whose inner bound is the outer parameter -- is summed, not multiplied. if isinstance(entry, nd.MapEntry): - nmap: nd.Map = entry.map - range: Range = nmap.range - n_exec = range.num_elements() - work = sp.simplify(work * n_exec.subs(equality_subs[0]).subs(equality_subs[1]).subs(subs1)) + for param, (begin, end, step) in zip(entry.map.params, entry.map.range): + work = accumulate_over_range(work, pystr_to_symbolic(param), begin, end, step, equality_subs, subs1) + work = simplify(work) else: - print('WARNING: Only Map scopes are supported in work analysis for now. Assuming 1 iteration.') + warnings.warn('Only Map scopes are supported in work analysis; assuming 1 iteration.') # Work inside a state can simply be summed up. But now we need to find the depth of a state (i.e. longest path). # Since dataflow graph is a DAG, this can be done in linear time. - max_depth = sp.sympify(0) + max_depth = pystr_to_symbolic(0) # only do this if we are analyzing depth if analyze_tasklet == get_tasklet_work_depth or analyze_tasklet == get_tasklet_avg_par: # Calculate the maximum depth of the scope by finding the 'deepest' path from the source to the sink. This is done by @@ -795,12 +1207,12 @@ def scope_work_depth( # find all starting nodes if entry: # the entry is the starting node - traversal_q.append((entry, sp.sympify(0), None)) + traversal_q.append((entry, pystr_to_symbolic(0), None)) else: for node in scope_nodes: if len(state.in_edges(node)) == 0: # This node is a start node of the traversal - traversal_q.append((node, sp.sympify(0), None)) + traversal_q.append((node, pystr_to_symbolic(0), None)) # this map keeps track of the length of the longest path ending at each state so far seen. depth_map = {} wcr_depth_map = {} @@ -810,7 +1222,7 @@ def scope_work_depth( if in_edge is not None: visited.add(in_edge) - n_depth = sp.simplify(in_depth + w_d_map[get_uuid(node, state)][1]) + n_depth = simplify(in_depth + w_d_map[get_uuid(node, state)][1]) if node in depth_map: depth_map[node] = sp.Max(depth_map[node], n_depth) @@ -829,7 +1241,7 @@ def scope_work_depth( out_edges = state.out_edges(exit_node) for oedge in out_edges: # check for wcr - wcr_depth = sp.sympify(0) + wcr_depth = pystr_to_symbolic(0) if oedge.data.wcr is not None: # This division gives us the number of writes to each single memory location, which is the depth # as these need to be sequential (without assumptions on HW etc). @@ -841,7 +1253,7 @@ def scope_work_depth( else: wcr_depth_map[get_uuid(node, state)] = wcr_depth # We do not need to propagate the wcr_depth to MapExits, since else this will result in depth N + 1 for Maps of range N. - wcr_depth = wcr_depth if not isinstance(oedge.dst, nd.MapExit) else sp.sympify(0) + wcr_depth = wcr_depth if not isinstance(oedge.dst, nd.MapExit) else pystr_to_symbolic(0) # only append if it's actually new information # this e.g. helps for huge nested SDFGs with lots of inputs/outputs inside a map scope @@ -896,10 +1308,13 @@ def analyze_sdfg(sdfg: SDFG, w_d_map: Dict[str, sp.Expr], analyze_tasklet, assumptions: List[str], - detailed_analysis: bool = False) -> None: + detailed_analysis: bool = False): """ Analyze a given SDFG. We can either analyze work, work and depth or average parallelism. + :note: Only structured control flow is supported (loops as ``LoopRegion``, branches as + ``ConditionalBlock``, no ``break`` / ``continue`` / ``return``). An SDFG with unstructured + control flow is not analyzed: the analysis warns and returns a zero result. :note: SDFGs should have split interstate edges. This means there should be no interstate edges containing both a condition and an assignment. :param sdfg: The SDFG to analyze. @@ -910,14 +1325,24 @@ def analyze_sdfg(sdfg: SDFG, and work depth values for both branches. If False, the worst-case branch is taken. Discouraged to use on bigger SDFGs, as computation time sky-rockets, since expression can became HUGE (depending on number of branches etc.). """ - # deepcopy such that original sdfg not changed sdfg = deepcopy(sdfg) + # The analysis only models structured control flow. If the SDFG has a legacy loop or + # unstructured branching, bail out with a zero result rather than producing a wrong one. + if has_unstructured_control_flow(sdfg): + warnings.warn('Work-depth analysis supports only structured control flow (LoopRegion / ' + 'ConditionalBlock); the SDFG contains a legacy loop or unstructured branch, ' + 'so no result is produced.') + result = (pystr_to_symbolic(0), + pystr_to_symbolic(0)) if analyze_tasklet == get_tasklet_work_depth else pystr_to_symbolic(0) + w_d_map[get_uuid(sdfg)] = result + return result + # apply SSA pass pipeline = FixedPointPipeline([StrictSymbolSSA()]) pipeline.apply_pass(sdfg, {}) - + static_symbol_mapping = get_static_symbols(sdfg) array_symbols = get_array_size_symbols(sdfg) # parse assumptions equality_subs, all_subs = parse_assumptions(assumptions if assumptions is not None else [], array_symbols) @@ -929,8 +1354,8 @@ def analyze_sdfg(sdfg: SDFG, # Analyze the work and depth of the SDFG. symbols = {} - sdfg_work_depth(sdfg, w_d_map, analyze_tasklet, symbols, equality_subs, all_subs[0][0] if len(all_subs) > 0 else {}, - detailed_analysis) + control_flow_region_work_depth(sdfg, w_d_map, analyze_tasklet, symbols, equality_subs, + all_subs[0][0] if len(all_subs) > 0 else {}, detailed_analysis) for k, (v_w, v_d) in w_d_map.items(): # The symeval replaces nested SDFG symbols with their global counterparts. @@ -939,6 +1364,24 @@ def analyze_sdfg(sdfg: SDFG, v_d = symeval(v_d, symbols) w_d_map[k] = (v_w, v_d) + for k, v, in w_d_map.items(): + w_d_map[k] = ((v[0].subs(static_symbol_mapping).subs(equality_subs[1])), + (v[1].subs(static_symbol_mapping).subs(equality_subs[1]))) + + if analyze_tasklet == get_tasklet_work_depth: + for k, v, in w_d_map.items(): + w_d_map[k] = ((simplify(v[0])), (simplify(v[1]))) + elif analyze_tasklet == get_tasklet_work: + for k, v, in w_d_map.items(): + w_d_map[k] = (simplify(v[0])) + elif analyze_tasklet == get_tasklet_avg_par: + for k, v, in w_d_map.items(): + w_d_map[k] = (simplify(v[0] / v[1]) if (v[1]) != 0 else 0) # work / depth = avg par + + result_whole_sdfg = w_d_map[get_uuid(sdfg)] + + return result_whole_sdfg + def do_subs(work, depth, all_subs): """ @@ -951,11 +1394,11 @@ def do_subs(work, depth, all_subs): # first do subs2 of first sub # then do all the remaining subs subs2 = all_subs[0][1] if len(all_subs) > 0 else {} - work, depth = sp.simplify(sp.sympify(work).subs(subs2)), sp.simplify(sp.sympify(depth).subs(subs2)) + work, depth = simplify(pystr_to_symbolic(work).subs(subs2)), simplify(pystr_to_symbolic(depth).subs(subs2)) for i in range(1, len(all_subs)): subs1, subs2 = all_subs[i] - work, depth = sp.simplify(work.subs(subs1)), sp.simplify(depth.subs(subs1)) - work, depth = sp.simplify(work.subs(subs2)), sp.simplify(depth.subs(subs2)) + work, depth = simplify(work.subs(subs1)), simplify(depth.subs(subs1)) + work, depth = simplify(work.subs(subs2)), simplify(depth.subs(subs2)) return work, depth @@ -995,16 +1438,6 @@ def main() -> None: work_depth_map = {} analyze_sdfg(sdfg, work_depth_map, analyze_tasklet, args.assume, args.detailed) - if args.analyze == 'workDepth': - for k, v, in work_depth_map.items(): - work_depth_map[k] = (str(sp.simplify(v[0])), str(sp.simplify(v[1]))) - elif args.analyze == 'work': - for k, v, in work_depth_map.items(): - work_depth_map[k] = str(sp.simplify(v[0])) - elif args.analyze == 'avgPar': - for k, v, in work_depth_map.items(): - work_depth_map[k] = str(sp.simplify(v[0] / v[1]) if str(v[1]) != '0' else 0) # work / depth = avg par - result_whole_sdfg = work_depth_map[get_uuid(sdfg)] print(80 * '-') diff --git a/dace/subsets.py b/dace/subsets.py index 5c40f83222..802ecb9545 100644 --- a/dace/subsets.py +++ b/dace/subsets.py @@ -1,4 +1,4 @@ -# Copyright 2019-2025 ETH Zurich and the DaCe authors. All rights reserved. +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. import dace.serialize from dace import symbolic import sympy as sp @@ -769,9 +769,20 @@ def __getitem__(self, key): def __setitem__(self, key, value): # ``__init__`` coerces every bound; this path did not, so ``r[i] = (0, n - 1, 1)`` # quietly put Python ints into a container whose contract says symbolic. + def coerce(idx, v): + if isinstance(v, (tuple, list)): + v = symbolic_range_tuple(v) + if len(v) == 4: + self.tile_sizes[idx] = v[3] + return v[:3] + # Single-index write (e.g. the frontend replacing one dimension by an + # expression): still coerce so no raw Python number slips in. + return symbolic.pystr_to_symbolic(v) + if isinstance(key, slice): - return self.ranges.__setitem__(key, [symbolic_range_tuple(v) for v in value]) - return self.ranges.__setitem__(key, symbolic_range_tuple(value)) + indices = range(*key.indices(len(self.ranges))) + return self.ranges.__setitem__(key, [coerce(i, v) for i, v in zip(indices, value)]) + return self.ranges.__setitem__(key, coerce(key, value)) def __eq__(self, other): if not isinstance(other, Range): diff --git a/dace/symbolic.py b/dace/symbolic.py index 282d7e5138..06b4cc6421 100644 --- a/dace/symbolic.py +++ b/dace/symbolic.py @@ -1,4 +1,4 @@ -# Copyright 2019-2025 ETH Zurich and the DaCe authors. All rights reserved. +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. import ast import contextlib from collections import Counter @@ -1064,6 +1064,8 @@ def eval(cls, x, y): if y.is_Number: if y == 1: return x + if y == -1: + return -x # Exact division is not a rounding operation at all -- return the quotient itself, so the # expression stays comparable and simplifiable instead of hiding behind an int_floor node. quotient = x / y @@ -1705,6 +1707,20 @@ def visit_Attribute(self, node): return ast.copy_location(new_node, node) +def _construct_function_uncached(func, *args, **kwargs): + # Construct without SymPy's ``@cacheit`` constructor caches (both + # ``Function.__new__`` and ``Application.__new__`` are cached, and ``eval`` + # implementations re-enter them): DaCe symbol equality ignores dtype, so a + # cache entry built from an equal-named, different-dtype symbol would + # silently substitute that symbol into the result. Symbol-free arguments + # hash soundly and keep the regular (evaluating) constructors. + if (isinstance(func, type) and issubclass(func, sympy.core.function.Application) + and not (set(kwargs) - {'evaluate'}) and not kwargs.get('evaluate', False) + and any(isinstance(arg, sympy.Basic) and arg.free_symbols for arg in args)): + return sympy.Basic.__new__(func, *args) + return func(*args, **kwargs) + + class _SerializedSymbolicParser(ast.NodeVisitor): """ Parser for the deterministic expression strings produced by @@ -1808,7 +1824,7 @@ def _binop_pow(a, b): @staticmethod def _binop_mod(a, b): - return sympy.Mod(a, b, evaluate=False) + return _construct_function_uncached(sympy.Mod, a, b, evaluate=False) @staticmethod def _unary_minus(a): @@ -1821,13 +1837,13 @@ def _unary_minus(a): ast.Div: _binop_div, ast.Pow: _binop_pow, ast.Mod: _binop_mod, - ast.FloorDiv: lambda a, b: int_floor(a, b), + ast.FloorDiv: lambda a, b: _construct_function_uncached(int_floor, a, b), } _unaryops = { ast.UAdd: lambda a: +a, ast.USub: _unary_minus, - ast.Not: lambda a: sympy.Not(a), - ast.Invert: lambda a: bitwise_invert(a), + ast.Not: lambda a: _construct_function_uncached(sympy.Not, a), + ast.Invert: lambda a: _construct_function_uncached(bitwise_invert, a), } _comparators = { ast.Eq: sympy.Eq, @@ -1929,20 +1945,22 @@ def visit_BoolOp(self, node): if isinstance(node.op, ast.And): result = values[0] for value in values[1:]: - result = AND(result, value) + result = _construct_function_uncached(AND, result, value) return result result = values[0] for value in values[1:]: - result = OR(result, value) + result = _construct_function_uncached(OR, result, value) return result def visit_Compare(self, node): if len(node.ops) != 1 or len(node.comparators) != 1: raise NotImplementedError('Chained comparisons are not supported in symbolic deserialization') - return self._comparators[type(node.ops[0])](self.visit(node.left), self.visit(node.comparators[0])) + return _construct_function_uncached(self._comparators[type(node.ops[0])], self.visit(node.left), + self.visit(node.comparators[0])) def visit_IfExp(self, node): - return IfExpr(self.visit(node.test), self.visit(node.body), self.visit(node.orelse)) + return _construct_function_uncached(IfExpr, self.visit(node.test), self.visit(node.body), + self.visit(node.orelse)) def visit_Call(self, node): if isinstance(node.func, ast.Name) and node.func.id == '__dace_typed_const__': @@ -1992,9 +2010,7 @@ def visit_Call(self, node): return _cast_symbolic_value(args[0], func) kwargs = {kw.arg: self.visit(kw.value) for kw in node.keywords} - if kwargs: - return func(*args, **kwargs) - return func(*args) + return _construct_function_uncached(func, *args, **kwargs) def visit_Attribute(self, node): if isinstance(node.value, ast.Name) and node.value.id == 'dace': @@ -2002,7 +2018,7 @@ def visit_Attribute(self, node): return getattr(dtypes, node.attr) except AttributeError as ex: raise TypeError(f'Unknown DaCe dtype "{node.attr}"') from ex - return Attr(self.visit(node.value), symbol(node.attr)) + return _construct_function_uncached(Attr, self.visit(node.value), symbol(node.attr)) def generic_visit(self, node): raise TypeError(f'Unsupported node in symbolic deserialization: {type(node).__name__}') @@ -2030,7 +2046,7 @@ def _cast_symbolic_value(value, dtype: dtypes.typeclass): return TypedConstant(value, dtype=dtype) # Non-constant composite expressions are preserved as explicit casts so they # can round-trip even when no constant/symbol dtype rewrite is possible. - return sympy.Function(f'dace.{dtype.to_string()}')(value) + return _construct_function_uncached(sympy.Function(f'dace.{dtype.to_string()}'), value) class DaceSympySerializer(sympy.printing.str.StrPrinter): diff --git a/tests/gpu_worker_pinning_test.py b/tests/gpu_worker_pinning_test.py new file mode 100644 index 0000000000..bc99bc5017 --- /dev/null +++ b/tests/gpu_worker_pinning_test.py @@ -0,0 +1,72 @@ +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. +""" +Unit tests for the root conftest.py GPU worker/rank pinning logic. +""" +import importlib.util +import os +import pathlib +import types + +from dace.sdfg.sdfg import LAUNCHER_RANK_VARS + +CONFTEST_PATH = pathlib.Path(__file__).resolve().parent.parent / 'conftest.py' + + +def load_root_conftest() -> types.ModuleType: + spec = importlib.util.spec_from_file_location('dace_root_conftest_under_test', CONFTEST_PATH) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +root_conftest = load_root_conftest() + + +def clear_worker_env(monkeypatch) -> None: + monkeypatch.delenv('PYTEST_XDIST_WORKER', raising=False) + for var in LAUNCHER_RANK_VARS: + monkeypatch.delenv(var, raising=False) + monkeypatch.delenv('CUDA_VISIBLE_DEVICES', raising=False) + + +def test_parse_worker_index_gw_prefix(): + assert root_conftest.parse_worker_index('gw7') == 7 + + +def test_parse_worker_index_fallback_no_digits(): + assert root_conftest.parse_worker_index('master') == 0 + assert root_conftest.parse_worker_index('gw') == 0 + + +def test_pick_gpu_worker_device_modulo_wrap(): + assert root_conftest.pick_gpu_worker_device('gw4', ['0', '1']) == '0' + assert root_conftest.pick_gpu_worker_device('gw5', ['0', '1']) == '1' + + +def test_resolve_worker_id_prefers_xdist_over_mpi_rank(monkeypatch): + clear_worker_env(monkeypatch) + monkeypatch.setenv('PYTEST_XDIST_WORKER', 'gw1') + monkeypatch.setenv('OMPI_COMM_WORLD_RANK', '5') + assert root_conftest.resolve_worker_id() == 'gw1' + + +def test_resolve_worker_id_uses_first_set_rank_var_in_order(monkeypatch): + clear_worker_env(monkeypatch) + monkeypatch.setenv('PMI_RANK', '9') + monkeypatch.setenv('SLURM_PROCID', '1') + assert root_conftest.resolve_worker_id() == '9' + + +def test_pin_worker_to_gpu_uses_preset_pool_and_rank_var(monkeypatch): + clear_worker_env(monkeypatch) + monkeypatch.setenv('OMPI_COMM_WORLD_RANK', '1') + monkeypatch.setenv('CUDA_VISIBLE_DEVICES', '2,3') + root_conftest.pin_worker_to_gpu() + assert os.environ['CUDA_VISIBLE_DEVICES'] == '3' + + +def test_pin_worker_to_gpu_noop_without_worker_id(monkeypatch): + clear_worker_env(monkeypatch) + monkeypatch.setenv('CUDA_VISIBLE_DEVICES', '0,1') + root_conftest.pin_worker_to_gpu() + assert os.environ['CUDA_VISIBLE_DEVICES'] == '0,1' diff --git a/tests/sdfg/operational_intensity_test.py b/tests/sdfg/operational_intensity_test.py index 36be0455db..42b92719e6 100644 --- a/tests/sdfg/operational_intensity_test.py +++ b/tests/sdfg/operational_intensity_test.py @@ -1,65 +1,71 @@ -# Copyright 2019-2023 ETH Zurich and the DaCe authors. All rights reserved. +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. """ Contains test cases for the operational intensity analysis. """ +import contextlib +import io from typing import Dict, Tuple +from unittest import mock import pytest -import dace as dc +import dace import sympy as sp import numpy as np from dace.sdfg.performance_evaluation.operational_intensity import analyze_sdfg_op_in from dace.sdfg.performance_evaluation.helpers import get_uuid +from dace.sdfg.utils import inline_control_flow_regions +from dace.symbolic import pystr_to_symbolic, SymbolicType from dace.frontend.python.parser import DaceProgram from math import isclose -N = dc.symbol('N') -M = dc.symbol('M') -K = dc.symbol('K') +N = dace.symbol('N') +M = dace.symbol('M') +K = dace.symbol('K') -TILE_SIZE = dc.symbol('TILE_SIZE') +TILE_SIZE = dace.symbol('TILE_SIZE') -@dc.program -def single_map64(x: dc.float64[N], y: dc.float64[N], z: dc.float64[N]): +@dace.program +def single_map64(x: dace.float64[N], y: dace.float64[N], z: dace.float64[N]): z[:] = x + y # does N work, loads 3*N elements of 8 bytes # --> op_in should be N / 3*8*N = 1/24 (no reuse) assuming L divides N -@dc.program -def single_map16(x: dc.float16[N], y: dc.float16[N], z: dc.float16[N]): +@dace.program +def single_map16(x: dace.float16[N], y: dace.float16[N], z: dace.float16[N]): z[:] = x + y # does N work, loads 3*N elements of 2 bytes # --> op_in should be N / 3*2*N = 1/6 (no reuse) assuming L divides N -@dc.program -def single_for_loop(x: dc.float64[N], y: dc.float64[N]): +@dace.program +def single_for_loop(x: dace.float64[N], y: dace.float64[N]): for i in range(N): x[i] += y[i] # N work, 2*N*8 bytes loaded # --> 1/16 op in -@dc.program -def if_else(x: dc.int64[100], sum: dc.int64[1]): +@dace.program +def if_else(x: dace.int64[100], sum: dace.int64[1]): if x[10] > 50: for i in range(100): sum += x[i] if x[0] > 3: for i in range(100): sum += x[i] - # no else --> simply analyze the ifs. if cache big enough, everything is reused + # no else --> simply analyze the ifs. if cache big enough, everything is reused; -@dc.program -def unaligned_for_loop(x: dc.float32[100], sum: dc.int64[1]): +@dace.program +def unaligned_for_loop(x: dace.float32[100], sum: dace.int64[1]): for i in range(17, 53): sum += x[i] + # 36 = 144byte array elemets accessed 1 = 4byte scalar accessed 36 ops -> 64byte line size=> 3 lines + 1 line scalar => op in = 9/64 -@dc.program -def sequential_maps(x: dc.float64[N], y: dc.float64[N], z: dc.float64[N]): +@dace.program +def sequential_maps(x: dace.float64[N], y: dace.float64[N], z: dace.float64[N]): z[:] = x + y z[:] *= 2 z[:] += x @@ -67,8 +73,8 @@ def sequential_maps(x: dc.float64[N], y: dc.float64[N], z: dc.float64[N]): # --> op_in should be N / 3*8*N = 1/24 (no reuse) assuming L divides N -@dc.program -def nested_reuse(x: dc.float64[N], y: dc.float64[N], z: dc.float64[N], result: dc.float64[1]): +@dace.program +def nested_reuse(x: dace.float64[N], y: dace.float64[N], z: dace.float64[N], result: dace.float64[1]): # load x, y and z z[:] = x + y result[0] = np.sum(z) @@ -76,39 +82,48 @@ def nested_reuse(x: dc.float64[N], y: dc.float64[N], z: dc.float64[N], result: d # to z outside of the nested SDFG. -@dc.program -def mmm(x: dc.float64[N, N], y: dc.float64[N, N], z: dc.float64[N, N]): - for n, k, m in dc.map[0:N, 0:N, 0:N]: +@dace.program +def mmm(x: dace.float64[N, N], y: dace.float64[N, N], z: dace.float64[N, N]): + for n, k, m in dace.map[0:N, 0:N, 0:N]: z[n, k] += x[n, m] * y[m, k] -@dc.program -def tiled_mmm(x: dc.float64[N, N], y: dc.float64[N, N], z: dc.float64[N, N]): - for n_TILE, k_TILE, m_TILE in dc.map[0:N:TILE_SIZE, 0:N:TILE_SIZE, 0:N:TILE_SIZE]: - for n, k, m in dc.map[n_TILE:n_TILE + TILE_SIZE, k_TILE:k_TILE + TILE_SIZE, m_TILE:m_TILE + TILE_SIZE]: +@dace.program +def tiled_mmm(x: dace.float64[N, N], y: dace.float64[N, N], z: dace.float64[N, N]): + for n_TILE, k_TILE, m_TILE in dace.map[0:N:TILE_SIZE, 0:N:TILE_SIZE, 0:N:TILE_SIZE]: + for n, k, m in dace.map[n_TILE:n_TILE + TILE_SIZE, k_TILE:k_TILE + TILE_SIZE, m_TILE:m_TILE + TILE_SIZE]: z[n, k] += x[n, m] * y[m, k] -@dc.program -def tiled_mmm_32(x: dc.float32[N, N], y: dc.float32[N, N], z: dc.float32[N, N]): - for n_TILE, k_TILE, m_TILE in dc.map[0:N:TILE_SIZE, 0:N:TILE_SIZE, 0:N:TILE_SIZE]: - for n, k, m in dc.map[n_TILE:n_TILE + TILE_SIZE, k_TILE:k_TILE + TILE_SIZE, m_TILE:m_TILE + TILE_SIZE]: +@dace.program +def tiled_mmm_32(x: dace.float32[N, N], y: dace.float32[N, N], z: dace.float32[N, N]): + for n_TILE, k_TILE, m_TILE in dace.map[0:N:TILE_SIZE, 0:N:TILE_SIZE, 0:N:TILE_SIZE]: + for n, k, m in dace.map[n_TILE:n_TILE + TILE_SIZE, k_TILE:k_TILE + TILE_SIZE, m_TILE:m_TILE + TILE_SIZE]: z[n, k] += x[n, m] * y[m, k] -@dc.program -def reduction_library_node(x: dc.float64[N]): +@dace.program +def reduction_library_node(x: dace.float64[N]): return np.sum(x) #(sdfg, c, l, assumptions, expected_result) -test_cases: Dict[str, Tuple[DaceProgram, int, int, Dict[str, int], dc.symbolic.SymbolicType]] = { +test_cases: Dict[str, Tuple[DaceProgram, int, int, Dict[str, int], SymbolicType]] = { 'single_map64_even': (single_map64, 64 * 64, 64, { 'N': 512 }, 1 / 24), 'single_map16_even': (single_map16, 64 * 64, 64, { 'N': 512 }, 1 / 6), + 'single_for_loop': (single_for_loop, 64 * 64, 64, { + 'N': 512 + }, 1 / 16), + 'if_else': (if_else, 64 * 64, 64, { + 'N': 512 + }, 200 / (14 * 64)) + # 14 cache misses, because DaCe introduces intermediate variable + , + 'unaligned_for_loop': (unaligned_for_loop, 64 * 64, 64, {}, 9 / 64), # now num_elements_on_single_cache_line does not divie N anymore # -->513 work, 520 elements loaded --> 513 / (520*8*3) 'single_map64_uneven': (single_map64, 64 * 64, 64, { @@ -137,7 +152,7 @@ def reduction_library_node(x: dc.float64[N]): }, (2 * 24**3) / (16 * 12 * 6**3)), 'reduction_library_node': (reduction_library_node, 1024, 64, { 'N': 128 - }, 128.0 / (dc.symbol('Reduce_misses') * 64.0 + 64.0)), + }, 128.0 / (dace.symbol('Reduce_misses', positive=True) * 64.0 + 64.0)), } @@ -153,17 +168,105 @@ def test_operational_intensity(test_name: str): analyze_sdfg_op_in(sdfg, op_in_map, c * l, l, assumptions) res = (op_in_map[get_uuid(sdfg)]) if test_name == 'reduction_library_node': - # substitue each symbol without assumptions. - # We do this since sp.Symbol('N') == Sp.Symbol('N', positive=True) --> False. - reps = {s: sp.Symbol(s.name) for s in res.free_symbols} - res = res.subs(reps) - reps = {s: sp.Symbol(s.name) for s in sp.sympify(correct).free_symbols} - correct = sp.sympify(correct).subs(reps) - assert correct == res + # Symbolic result (depends on the opaque Reduce_misses symbol); compare expressions directly. + assert pystr_to_symbolic(correct) == res else: assert isclose(correct, res) +_ASK_USER_LOOP_ITERS = 8 + + +@dace.program +def ask_user_branch(x: dace.float64[64], y: dace.float64[64]): + # Data-dependent branches doing different amounts of work, so the chosen branch changes the result. + if x[0] > 0: + y[:] = x + 1.0 + else: + y[:] = x * x * x * x + + +@dace.program +def ask_user_branch_in_loop(x: dace.float64[64], y: dace.float64[64]): + for _ in range(_ASK_USER_LOOP_ITERS): + if x[0] > 0: + y[:] = x + 1.0 + else: + y[:] = x * x * x * x + + +def _op_in_with_choice(program: DaceProgram, choice: int) -> Tuple[int, float]: + """ Run the analysis in ``ask_user`` mode, answering every branch prompt with ``choice``. + + :param program: The DaCe program to analyze. + :param choice: The branch index to feed to every prompt. + :returns: The number of prompts raised and the resulting operational intensity. + """ + sdfg = program.to_sdfg() + sdfg.simplify() + prompts = [] + + def fake_input(*_): + prompts.append(choice) + return str(choice) + + op_in_map: Dict[str, sp.Expr] = {} + with mock.patch('builtins.input', fake_input), contextlib.redirect_stdout(io.StringIO()): + analyze_sdfg_op_in(sdfg, op_in_map, 1024, 64, {}, ask_user=True) + return len(prompts), float(op_in_map[get_uuid(sdfg)]) + + +def test_operational_intensity_ask_user_branch_selection(): + """ ``ask_user`` picks which branch of a data-dependent conditional to analyze. Both choices + must complete without error and, since the branches differ in work, yield different results. """ + prompts_true, op_in_true = _op_in_with_choice(ask_user_branch, 0) + prompts_else, op_in_else = _op_in_with_choice(ask_user_branch, 1) + assert prompts_true == 1 and prompts_else == 1 + assert not isclose(op_in_true, op_in_else) + + +def test_operational_intensity_ask_user_decision_reused_in_loop(): + """ A branch chosen once is reused on later visits, so a conditional inside a loop prompts only + once and the loop intensity is the single-iteration intensity scaled by the trip count. """ + single_prompts, single_op_in = _op_in_with_choice(ask_user_branch, 0) + loop_prompts, loop_op_in = _op_in_with_choice(ask_user_branch_in_loop, 0) + assert single_prompts == 1 and loop_prompts == 1 + assert isclose(loop_op_in, _ASK_USER_LOOP_ITERS * single_op_in) + + +def test_operational_intensity_range_simulation(): + """ Smoke-test the simulation path: a symbol given a range ``'start,stop,step'`` is sampled, its + cache misses simulated per sample, and the operational intensity fitted as a function of it. + Streaming ``single_map64`` has no reuse, so the fit is the constant ``1 / 24``. """ + op_in_map: Dict[str, sp.Expr] = {} + sdfg = single_map64.to_sdfg() + # Sampling at multiples of 8 keeps the 64-byte cache lines (8 doubles) evenly divided. + analyze_sdfg_op_in(sdfg, op_in_map, 64 * 64, 64, {'N': '64,576,64'}, test_set_size=2) + + # The fitted result is a string expression in N; parse it (``pystr_to_symbolic`` avoids the + # collision between the symbol ``N`` and SymPy's numeric-evaluation function ``N``). + op_in = pystr_to_symbolic(op_in_map[get_uuid(sdfg)]) + for n in (64, 256, 512): + assert isclose(float(op_in.subs(N, n)), 1 / 24, rel_tol=1e-6) + + +def test_operational_intensity_bails_on_unstructured_control_flow(): + """ The analysis only models structured control flow. On an inlined SDFG (LoopRegions and + ConditionalBlocks flattened to a legacy state machine) it must warn and produce a zero result + rather than a wrong one. """ + sdfg = ask_user_branch.to_sdfg() + sdfg.simplify() + inline_control_flow_regions(sdfg) + op_in_map: Dict[str, sp.Expr] = {} + with pytest.warns(UserWarning, match='structured control flow'): + analyze_sdfg_op_in(sdfg, op_in_map, 1024, 64, {}) + assert op_in_map[get_uuid(sdfg)] == 0 + + if __name__ == '__main__': for test_name in test_cases.keys(): test_operational_intensity(test_name) + test_operational_intensity_ask_user_branch_selection() + test_operational_intensity_ask_user_decision_reused_in_loop() + test_operational_intensity_range_simulation() + test_operational_intensity_bails_on_unstructured_control_flow() diff --git a/tests/sdfg/work_depth_test.py b/tests/sdfg/work_depth_test.py index d2323aeddd..515c7c994d 100644 --- a/tests/sdfg/work_depth_test.py +++ b/tests/sdfg/work_depth_test.py @@ -1,10 +1,10 @@ -# Copyright 2019-2024 ETH Zurich and the DaCe authors. All rights reserved. +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. """ Contains test cases for the work depth analysis. """ from typing import Dict, List, Tuple import pytest -import dace as dc -from dace import symbolic +import dace +from dace.symbolic import pystr_to_symbolic, simplify, SymbolicType from dace.frontend.python.parser import DaceProgram from dace.sdfg.performance_evaluation.work_depth import (analyze_sdfg, get_tasklet_work_depth, get_tasklet_avg_par, parse_assumptions, count_arithmetic_ops_code, count_depth_code) @@ -19,24 +19,24 @@ from pytest import raises -N = dc.symbol('N') -M = dc.symbol('M') -K = dc.symbol('K') +N = dace.symbol('N') +M = dace.symbol('M') +K = dace.symbol('K') -@dc.program -def single_map(x: dc.float64[N], y: dc.float64[N], z: dc.float64[N]): +@dace.program +def single_map(x: dace.float64[N], y: dace.float64[N], z: dace.float64[N]): z[:] = x + y -@dc.program -def single_for_loop(x: dc.float64[N], y: dc.float64[N]): +@dace.program +def single_for_loop(x: dace.float64[N], y: dace.float64[N]): for i in range(N): x[i] += y[i] -@dc.program -def if_else(x: dc.int64[1000], y: dc.int64[1000], z: dc.int64[1000], sum: dc.int64[1]): +@dace.program +def if_else(x: dace.int64[1000], y: dace.int64[1000], z: dace.int64[1000], sum: dace.int64[1]): if x[10] > 50: z[:] = x + y # 1000 work, 1 depth else: @@ -44,8 +44,8 @@ def if_else(x: dc.int64[1000], y: dc.int64[1000], z: dc.int64[1000], sum: dc.int sum += x[i] -@dc.program -def if_else_sym(x: dc.int64[N], y: dc.int64[N], z: dc.int64[N], sum: dc.int64[1]): +@dace.program +def if_else_sym(x: dace.int64[N], y: dace.int64[N], z: dace.int64[N], sum: dace.int64[1]): if x[10] > 50: z[:] = x + y # N work, 1 depth else: @@ -53,26 +53,26 @@ def if_else_sym(x: dc.int64[N], y: dc.int64[N], z: dc.int64[N], sum: dc.int64[1] sum += x[i] -@dc.program -def nested_sdfg(x: dc.float64[N], y: dc.float64[N], z: dc.float64[N]): +@dace.program +def nested_sdfg(x: dace.float64[N], y: dace.float64[N], z: dace.float64[N]): single_map(x, y, z) single_for_loop(x, y) -@dc.program -def nested_maps(x: dc.float64[N, M], y: dc.float64[N, M], z: dc.float64[N, M]): +@dace.program +def nested_maps(x: dace.float64[N, M], y: dace.float64[N, M], z: dace.float64[N, M]): z[:, :] = x + y -@dc.program -def nested_for_loops(x: dc.float64[N], y: dc.float64[K]): +@dace.program +def nested_for_loops(x: dace.float64[N], y: dace.float64[K]): for i in range(N): for j in range(K): x[i] += y[j] -@dc.program -def nested_if_else(x: dc.int64[N], y: dc.int64[N], z: dc.int64[N], sum: dc.int64[1]): +@dace.program +def nested_if_else(x: dace.int64[N], y: dace.int64[N], z: dace.int64[N], sum: dace.int64[1]): if x[10] > 50: if x[9] > 40: z[:] = x + y # N work, 1 depth @@ -89,8 +89,8 @@ def nested_if_else(x: dc.int64[N], y: dc.int64[N], z: dc.int64[N], sum: dc.int64 # --> total over both branches: Max(K, M+N, 3*N) work, Max(K, M+1, 3) depth -@dc.program -def max_of_positive_symbol(x: dc.float64[N]): +@dace.program +def max_of_positive_symbol(x: dace.float64[N]): if x[0] > 0: for i in range(2 * N): # work 2*N^2, depth 2*N x += 1 @@ -100,9 +100,9 @@ def max_of_positive_symbol(x: dc.float64[N]): # total is work 3*N^2, depth 3*N without any max -@dc.program -def multiple_array_sizes(x: dc.int64[N], y: dc.int64[N], z: dc.int64[N], x2: dc.int64[M], y2: dc.int64[M], - z2: dc.int64[M], x3: dc.int64[K], y3: dc.int64[K], z3: dc.int64[K]): +@dace.program +def multiple_array_sizes(x: dace.int64[N], y: dace.int64[N], z: dace.int64[N], x2: dace.int64[M], y2: dace.int64[M], + z2: dace.int64[M], x3: dace.int64[K], y3: dace.int64[K], z3: dace.int64[K]): if x[0] > 0: z[:] = 2 * x + y # work 2*N, depth 2 elif x[1] > 0: @@ -115,14 +115,14 @@ def multiple_array_sizes(x: dc.int64[N], y: dc.int64[N], z: dc.int64[N], x2: dc. # --> work= Max(3*N, 2*M, 2*K) and depth = 5 -@dc.program -def unbounded_while_do(x: dc.float64[N]): +@dace.program +def unbounded_while_do(x: dace.float64[N]): while x[0] < 100: x += 1 -@dc.program -def unbounded_nonnegify(x: dc.float64[N]): +@dace.program +def unbounded_nonnegify(x: dace.float64[N]): while x[0] < 100: if x[1] < 42: x += 3 * x @@ -130,24 +130,31 @@ def unbounded_nonnegify(x: dc.float64[N]): x += x -@dc.program -def break_for_loop(x: dc.float64[N]): +@dace.program +def break_for_loop(x: dace.float64[N]): for i in range(N): if x[i] > 100: break x += 1 -@dc.program -def break_while_loop(x: dc.float64[N]): +@dace.program +def break_while_loop(x: dace.float64[N]): while x[0] > 10: if x[1] > 100: break x += 1 -@dc.program -def sequntial_ifs(x: dc.float64[N + 1], y: dc.float64[M + 1]): # --> cannot assume N, M to be positive +@dace.program +def early_return(x: dace.float64[N]): + if x[0] > 0: + return + x += 1 + + +@dace.program +def sequntial_ifs(x: dace.float64[N + 1], y: dace.float64[M + 1]): # --> cannot assume N, M to be positive if x[0] > 5: x[:] += 1 # N+1 work, 1 depth else: @@ -161,28 +168,34 @@ def sequntial_ifs(x: dc.float64[N + 1], y: dc.float64[M + 1]): # --> cannot ass # Depth: Max(1, M) + 1 -@dc.program -def reduction_library_node(x: dc.float64[456]): +@dace.program +def reduction_library_node(x: dace.float64[456]): return np.sum(x) -@dc.program -def reduction_library_node_symbolic(x: dc.float64[N]): +@dace.program +def reduction_library_node_symbolic(x: dace.float64[N]): return np.sum(x) -@dc.program -def gemm_library_node(x: dc.float64[456, 200], y: dc.float64[200, 111], z: dc.float64[456, 111]): +@dace.program +def gemm_library_node(x: dace.float64[456, 200], y: dace.float64[200, 111], z: dace.float64[456, 111]): z[:] = x @ y -@dc.program -def gemm_library_node_symbolic(x: dc.float64[M, K], y: dc.float64[K, N], z: dc.float64[M, N]): +@dace.program +def gemm_library_node_symbolic(x: dace.float64[M, K], y: dace.float64[K, N], z: dace.float64[M, N]): z[:] = x @ y +@dace.program +def loop_var_dependent_work(x: dace.float64[N], y: dace.float64[N], z: dace.float64[N]): + for i in range(1, N + 1): + z[i - 1] = np.dot(x[:i], y[:i]) + + #(sdfg, (expected_work, expected_depth)) -work_depth_test_cases: Dict[str, Tuple[DaceProgram, Tuple[symbolic.SymbolicType, symbolic.SymbolicType]]] = { +work_depth_test_cases: Dict[str, Tuple[DaceProgram, Tuple[SymbolicType, SymbolicType]]] = { 'single_map': (single_map, (N, 1)), 'single_for_loop': (single_for_loop, (N, N)), 'if_else': (if_else, (1000, 100)), @@ -193,23 +206,36 @@ def gemm_library_node_symbolic(x: dc.float64[M, K], y: dc.float64[K, N], z: dc.f 'nested_if_else': (nested_if_else, (sp.Max(K, 3 * N, M + N), sp.Max(3, K, M + 1))), 'max_of_positive_symbols': (max_of_positive_symbol, (3 * N**2, 3 * N)), 'multiple_array_sizes': (multiple_array_sizes, (sp.Max(2 * K, 3 * N, 2 * M + 3), 5)), - 'unbounded_while_do': (unbounded_while_do, (sp.Symbol('num_execs_0_5') * N, sp.Symbol('num_execs_0_5'))), + 'unbounded_while_do': (unbounded_while_do, (dace.symbol('num_execs_0_0', nonnegative=True) * N, + dace.symbol('num_execs_0_0', nonnegative=True))), # We get this Max(1, num_execs), since it is a do-while loop, but the num_execs symbol does not capture this. - 'unbounded_nonnegify': (unbounded_nonnegify, (2 * sp.Symbol('num_execs_0_8') * N, 2 * sp.Symbol('num_execs_0_8'))), - 'break_for_loop': (break_for_loop, (N**2, N)), - 'break_while_loop': (break_while_loop, (sp.Symbol('num_execs_0_7') * N, sp.Symbol('num_execs_0_7'))), + 'unbounded_nonnegify': (unbounded_nonnegify, (2 * dace.symbol('num_execs_0_0', nonnegative=True) * N, + 2 * dace.symbol('num_execs_0_0', nonnegative=True))), 'sequential_ifs': (sequntial_ifs, (sp.Max(N + 1, M) + sp.Max(N + 1, M + 1), sp.Max(1, M) + 1)), - 'reduction_library_node': (reduction_library_node, (456, sp.log(456))), - 'reduction_library_node_symbolic': (reduction_library_node_symbolic, (N, sp.log(N))), - 'gemm_library_node': (gemm_library_node, (2 * 456 * 200 * 111, sp.log(200))), - 'gemm_library_node_symbolic': (gemm_library_node_symbolic, (2 * M * K * N, sp.log(K))) + 'reduction_library_node': (reduction_library_node, (456, sp.log(456) / sp.log(2))), + 'reduction_library_node_symbolic': (reduction_library_node_symbolic, (N, sp.log(sp.Max(1, N)) / sp.log(2))), + 'gemm_library_node': (gemm_library_node, (2 * 456 * 200 * 111, sp.log(200) / sp.log(2))), + 'gemm_library_node_symbolic': + (gemm_library_node_symbolic, (2 * M * K * N, sp.Max(1, + sp.log(sp.Max(1, K)) / sp.log(2)))), + 'loop_var_dependent_work': + (loop_var_dependent_work, (N**2, N + sp.Sum(sp.log(dace.symbol("_p_i", nonnegative=True) + 1), + (dace.symbol("_p_i", nonnegative=True), 0, N - 1)) / sp.log(2))) } +def assert_symbolically_equal(res: sp.Expr, correct: sp.Expr) -> None: + """ Assert that an analysis result is exactly the expected value, whatever shape sympy left it in. """ + # sympy.simplify is not idempotent on logs of composite integers (log(456)/log(2) and + # 3 + log(57)/log(2) map to each other), so the shape of a result depends on how many times the + # traversal simplified it, which in turn depends on the SDFG's state count. Compare values. + assert res.expand() == correct.expand() or simplify(res - correct) == 0 + + @pytest.mark.parametrize('test_name', list(work_depth_test_cases.keys())) def test_work_depth(test_name): - if (dc.Config.get_bool('optimizer', 'automatic_simplification') == False - and test_name in ['unbounded_while_do', 'unbounded_nonnegify', 'break_while_loop']): + if (dace.Config.get_bool('optimizer', 'automatic_simplification') == False + and test_name in ['unbounded_while_do', 'unbounded_nonnegify']): pytest.skip('Malformed loop when not simplifying') test, correct = work_depth_test_cases[test_name] w_d_map: Dict[str, sp.Expr] = {} @@ -219,21 +245,11 @@ def test_work_depth(test_name): if 'nested_maps' in test.name: sdfg.apply_transformations(MapExpansion) - # NOTE: Until the W/D Analysis is changed to make use of the new blocks, inline control flow for the analysis. - inline_control_flow_regions(sdfg) - for sd in sdfg.all_sdfgs_recursive(): - sd.using_explicit_control_flow = False - analyze_sdfg(sdfg, w_d_map, get_tasklet_work_depth, [], False) res = w_d_map[get_uuid(sdfg)] - # substitue each symbol without assumptions. - # We do this since sp.Symbol('N') == Sp.Symbol('N', positive=True) --> False. - reps = {s: sp.Symbol(s.name) for s in (res[0].free_symbols | res[1].free_symbols)} - res = (res[0].subs(reps), res[1].subs(reps)) - reps = {s: sp.Symbol(s.name) for s in (sp.sympify(correct[0]).free_symbols | sp.sympify(correct[1]).free_symbols)} - correct = (sp.sympify(correct[0]).subs(reps), sp.sympify(correct[1]).subs(reps)) - # check result - assert correct == res + correct = (pystr_to_symbolic(correct[0]), pystr_to_symbolic(correct[1])) + assert_symbolically_equal(res[0], correct[0]) + assert_symbolically_equal(res[1], correct[1]) #(sdfg, expected_avg_par) @@ -247,19 +263,19 @@ def test_work_depth(test_name): 'max_of_positive_symbol': (max_of_positive_symbol, N), 'unbounded_while_do': (unbounded_while_do, N), 'unbounded_nonnegify': (unbounded_nonnegify, N), - 'break_for_loop': (break_for_loop, N), - 'break_while_loop': (break_while_loop, N), - 'reduction_library_node': (reduction_library_node, 456 / sp.log(456)), - 'reduction_library_node_symbolic': (reduction_library_node_symbolic, N / sp.log(N)), - 'gemm_library_node': (gemm_library_node, 2 * 456 * 200 * 111 / sp.log(200)), - 'gemm_library_node_symbolic': (gemm_library_node_symbolic, 2 * M * K * N / sp.log(K)), + 'reduction_library_node': (reduction_library_node, 456 / (sp.log(456) / sp.log(2))), + 'reduction_library_node_symbolic': (reduction_library_node_symbolic, N * sp.log(2) / sp.log(sp.Max(1, N))), + 'gemm_library_node': (gemm_library_node, 2 * 456 * 200 * 111 / (sp.log(200) / sp.log(2))), + 'gemm_library_node_symbolic': + (gemm_library_node_symbolic, 2 * K * M * N / sp.Max(1, + sp.log(sp.Max(1, K)) / sp.log(2))), } @pytest.mark.parametrize('test_name', list(tests_cases_avg_par.keys())) def test_avg_par(test_name: str): - if (dc.Config.get_bool('optimizer', 'automatic_simplification') == False - and test_name in ['unbounded_while_do', 'unbounded_nonnegify', 'break_while_loop']): + if (dace.Config.get_bool('optimizer', 'automatic_simplification') == False + and test_name in ['unbounded_while_do', 'unbounded_nonnegify']): pytest.skip('Malformed loop when not simplifying') test, correct = tests_cases_avg_par[test_name] @@ -270,24 +286,37 @@ def test_avg_par(test_name: str): if 'nested_maps' in test_name: sdfg.apply_transformations(MapExpansion) - # NOTE: Until the W/D Analysis is changed to make use of the new blocks, inline control flow for the analysis. + analyze_sdfg(sdfg, w_d_map, get_tasklet_avg_par, [], False) + res = w_d_map[get_uuid(sdfg)] + correct = pystr_to_symbolic(correct) + assert_symbolically_equal(res, correct) + + +@pytest.mark.parametrize('prog', [break_for_loop, break_while_loop, early_return]) +def test_work_depth_bails_on_nonlocal_exit(prog: DaceProgram): + """ ``break`` / ``continue`` / ``return`` are not supported (non-local exits are not modeled); + the analysis must warn and produce a zero (work, depth) result rather than a wrong one. """ + sdfg = prog.to_sdfg() + w_d_map: Dict[str, sp.Expr] = {} + with pytest.warns(UserWarning, match='structured control flow'): + analyze_sdfg(sdfg, w_d_map, get_tasklet_work_depth, [], False) + assert w_d_map[get_uuid(sdfg)] == (0, 0) + + +def test_work_depth_bails_on_unstructured_control_flow(): + """ Inlined control flow (LoopRegions / ConditionalBlocks flattened to a legacy state machine) + is not supported; the analysis must warn and produce a zero (work, depth) result. """ + sdfg = single_for_loop.to_sdfg() inline_control_flow_regions(sdfg) for sd in sdfg.all_sdfgs_recursive(): sd.using_explicit_control_flow = False - - analyze_sdfg(sdfg, w_d_map, get_tasklet_avg_par, [], False) - res = w_d_map[get_uuid(sdfg)][0] / w_d_map[get_uuid(sdfg)][1] - # substitue each symbol without assumptions. - # We do this since sp.Symbol('N') == Sp.Symbol('N', positive=True) --> False. - reps = {s: sp.Symbol(s.name) for s in res.free_symbols} - res = res.subs(reps) - reps = {s: sp.Symbol(s.name) for s in sp.sympify(correct).free_symbols} - correct = sp.sympify(correct).subs(reps) - # check result - assert correct == res + w_d_map: Dict[str, sp.Expr] = {} + with pytest.warns(UserWarning, match='structured control flow'): + analyze_sdfg(sdfg, w_d_map, get_tasklet_work_depth, [], False) + assert w_d_map[get_uuid(sdfg)] == (0, 0) -x, y, z, a = sp.symbols('x y z a') +x, y, z, a = dace.symbol('x'), dace.symbol('y'), dace.symbol('z'), dace.symbol('a') # (expr, assumptions, result) assumptions_tests = [ @@ -331,7 +360,6 @@ def test_depth_counter_vs_work_counter(): """ Test that the DepthCounter correctly computes depth (longest chain of dependent operations) which can differ from work (total number of operations). - Depth measures the critical path through the expression tree, while work measures the total number of operations. """