Control Flow Block Constraints

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):
         """

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, {}

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()