diff --git a/dace/sdfg/state.py b/dace/sdfg/state.py index cafea3d754..429fbbd690 100644 --- a/dace/sdfg/state.py +++ b/dace/sdfg/state.py @@ -1092,6 +1092,10 @@ class ControlFlowBlock(BlockGraphView, abc.ABC): is_collapsed = Property(dtype=bool, desc='Show this block as collapsed', default=False) + pre_conditions = DictProperty(key_type=str, value_type=list, desc='Pre-conditions for this block') + post_conditions = DictProperty(key_type=str, value_type=list, desc='Post-conditions for this block') + invariant_conditions = DictProperty(key_type=str, value_type=list, desc='Invariant conditions for this block') + _label: str def __init__(self, @@ -1104,6 +1108,9 @@ def __init__(self, self._sdfg = sdfg self._parent_graph = parent self.is_collapsed = False + self.pre_conditions = {} + self.post_conditions = {} + self.invariant_conditions = {} def set_default_lineinfo(self, lineinfo: dace.dtypes.DebugInfo): """ diff --git a/dace/transformation/passes/analysis.py b/dace/transformation/passes/analysis.py index cccfbf10a3..82cae6e470 100644 --- a/dace/transformation/passes/analysis.py +++ b/dace/transformation/passes/analysis.py @@ -2,7 +2,7 @@ from collections import defaultdict from dace.transformation import pass_pipeline as ppl -from dace import SDFG, SDFGState, properties, InterstateEdge, Memlet, data as dt +from dace import SDFG, SDFGState, properties, InterstateEdge, Memlet, data as dt, symbolic from dace.sdfg.graph import Edge from dace.sdfg import nodes as nd from dace.sdfg.analysis import cfg @@ -583,3 +583,38 @@ def apply_pass(self, top_sdfg: SDFG, _) -> Dict[int, Dict[str, Set[Union[Memlet, result[anode.data].add(e.data) top_result[sdfg.cfg_id] = result return top_result + + +@properties.make_properties +class DeriveSDFGConstraints(ppl.Pass): + + CATEGORY: str = 'Analysis' + + assume_max_data_size = properties.Property(dtype=int, default=None, allow_none=True, + desc='Assume that all data containers have no dimension larger than ' + + 'this value. If None, no assumption is made.') + + def modifies(self) -> ppl.Modifies: + return ppl.Modifies.Nothing + + def should_reapply(self, modified: ppl.Modifies) -> bool: + # If anything was modified, reapply + return modified & ppl.Modifies.Everything + + def _derive_parameter_datasize_constraints(self, sdfg: SDFG, invariants: Dict[str, Set[str]]) -> None: + handled = set() + for arr in sdfg.arrays.values(): + for dim in arr.shape: + if isinstance(dim, symbolic.symbol) and not dim in handled: + ds = str(dim) + if ds not in invariants: + invariants[ds] = set() + invariants[ds].add(f'{ds} > 0') + if self.assume_max_data_size is not None: + invariants[ds].add(f'{ds} <= {self.assume_max_data_size}') + handled.add(ds) + + def apply_pass(self, sdfg: SDFG, _) -> Tuple[Dict[str, Set[str]], Dict[str, Set[str]], Dict[str, Set[str]]]: + invariants: Dict[str, Set[str]] = {} + self._derive_parameter_datasize_constraints(sdfg, invariants) + return {}, invariants, {} diff --git a/tests/passes/sdfg_constraint_derivation_test.py b/tests/passes/sdfg_constraint_derivation_test.py new file mode 100644 index 0000000000..868548da7f --- /dev/null +++ b/tests/passes/sdfg_constraint_derivation_test.py @@ -0,0 +1,49 @@ +# Copyright 2019-2024 ETH Zurich and the DaCe authors. All rights reserved. + +import dace +from dace.transformation.passes.analysis import DeriveSDFGConstraints + + +def test_infer_data_dim_constraints_nomax(): + N = dace.symbol('N') + + @dace.program + def matmul(A: dace.float64[N, N], B: dace.float64[N, N], C: dace.float64[N, N]): + for i in range(N): + for j in range(N): + for k in range(N): + C[i, j] += A[i, k] * B[k, j] + + sdfg = matmul.to_sdfg() + + derive_pass = DeriveSDFGConstraints() + _, inv, _ = derive_pass.apply_pass(sdfg, {}) + + assert 'N' in inv + assert 'N > 0' in inv['N'] + + +def test_infer_data_dim_constraints_withmax(): + N = dace.symbol('N') + + @dace.program + def matmul(A: dace.float64[N, N], B: dace.float64[N, N], C: dace.float64[N, N]): + for i in range(N): + for j in range(N): + for k in range(N): + C[i, j] += A[i, k] * B[k, j] + + sdfg = matmul.to_sdfg() + + derive_pass = DeriveSDFGConstraints() + derive_pass.assume_max_data_size = 128 + _, inv, _ = derive_pass.apply_pass(sdfg, {}) + + assert 'N' in inv + assert 'N > 0' in inv['N'] + assert 'N <= 128' in inv['N'] + + +if __name__ == "__main__": + test_infer_data_dim_constraints_nomax() + test_infer_data_dim_constraints_withmax()