Detect same-name symbols with differing SymPy assumptions

dace/frontend/python/newast.py

@@ -1,4 +1,4 @@
-# Copyright 2019-2024 ETH Zurich and the DaCe authors. All rights reserved.
+# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved.
 import ast
 from collections import OrderedDict
 import copy
@@ -28,6 +28,7 @@
 from dace.memlet import Memlet
 from dace.properties import LambdaProperty, CodeBlock
 from dace.sdfg import SDFG, SDFGState
+from dace.sdfg.sdfg import absorb_symbol_assumptions
 from dace.sdfg.state import (BreakBlock, ConditionalBlock, ContinueBlock, ControlFlowBlock, FunctionCallRegion,
                              LoopRegion, ControlFlowRegion, NamedRegion)
 from dace.sdfg.replace import replace_datadesc_names
@@ -1180,6 +1181,7 @@ def __init__(self,
                 for sym in arr.free_symbols:
                     if sym.name not in self.sdfg.symbols:
                         self.sdfg.add_symbol(sym.name, sym.dtype)
+            absorb_symbol_assumptions(self.sdfg, list(self.sdfg.arrays.values()))
         self.cfg_target = self.sdfg
         self.current_state = self.sdfg.add_state('init', is_start_block=True)
         self.last_block = self.current_state
@@ -2462,32 +2464,23 @@ def visit_For(self, node: ast.For):
         elif iterator == 'range':
             # Create an extra typed symbol for the loop iterate
             sym_name = indices[0]
-            integer = True
-            nonnegative = None
-            positive = None
+            # Mint the spelling a reparse produces. SymPy folds assumptions into symbol identity,
+            # so an iterator stamped `nonnegative=True` (or with an explicit `None`) is a DIFFERENT
+            # object from the `i` that every reparse of a loop bound or subset yields, and the two
+            # never cancel. Subset covering re-derives nonnegativity itself, via `subsets.nng`.
+            assumptions = {'integer': True}
 
             start = self._replace_with_global_symbols(symbolic.pystr_to_symbolic(ranges[0][0]))
             stop = self._replace_with_global_symbols(symbolic.pystr_to_symbolic(ranges[0][1]))
             step = self._replace_with_global_symbols(symbolic.pystr_to_symbolic(ranges[0][2]))
             eoff = -1
             if (step < 0) == True:
                 eoff = 1
-            try:
-                conditions = [s >= 0 for s in (start, stop, step)]
-                if (conditions == [True, True, True] or (start > stop and step < 0)):
-                    nonnegative = True
-                    if start != 0:
-                        positive = True
-            except:
-                pass
-
-            sym_obj = symbolic.symbol(indices[0],
-                                      dtypes.result_type_of(infer_expr_type(ranges[0][0], self.sdfg.symbols),
-                                                            infer_expr_type(ranges[0][1], self.sdfg.symbols),
-                                                            infer_expr_type(ranges[0][2], self.sdfg.symbols)),
-                                      integer=integer,
-                                      nonnegative=nonnegative,
-                                      positive=positive)
+            sym_obj = symbolic.symbol(
+                indices[0],
+                dtypes.result_type_of(infer_expr_type(ranges[0][0], self.sdfg.symbols),
+                                      infer_expr_type(ranges[0][1], self.sdfg.symbols),
+                                      infer_expr_type(ranges[0][2], self.sdfg.symbols)), **assumptions)
 
             if sym_name not in self.sdfg.symbols:
                 sym_name = self.sdfg.add_symbol(sym_name, sym_obj.dtype, find_new_name=True)

dace/sdfg/sdfg.py

@@ -11,6 +11,7 @@
 import random
 import shutil
 import sys
+import sympy
 from typing import Any, AnyStr, Dict, List, Optional, Sequence, Set, Tuple, Type, TYPE_CHECKING, Union
 import warnings
 
@@ -430,6 +431,36 @@ def from_json(json_obj, context=None):
         return ret
 
 
+def absorb_symbol_assumptions(sdfg: 'SDFG', descs: List[dt.Data]):
+    """Move assumptions off the symbols stored in ``descs`` into ``sdfg``'s registry.
+
+    Symbols stored in an SDFG are BARE: identity is the name. :class:`~dace.symbolic.UndefinedSymbol`
+    is skipped: it is a runtime-deferred sentinel spelled "?", it holds no assumptions, and its name
+    is not a legal symbol name, so it never enters the registry.
+
+    :param sdfg: The SDFG whose registry receives the facts.
+    :param descs: Data descriptors to normalize, rewritten in place.
+    :return: ``None``.
+    :rtype: None
+    """
+    assumed = {}
+    for desc in descs:
+        for sym in desc.free_symbols:
+            if isinstance(sym, symbolic.UndefinedSymbol) or sym in assumed:
+                continue
+            bare = symbolic.symbol(sym.name, sym.dtype)
+            known = bare.assumptions0
+            facts = {k: v for k, v in sym.assumptions0.items() if known.get(k) != v}
+            if not facts:
+                continue
+            if sym.name in sdfg.symbols:
+                sdfg.update_symbol_assumptions(sym.name, **facts)
+            assumed[sym] = bare
+    if assumed:
+        for desc in descs:
+            replace_properties_dict(desc, {}, assumed)
+
+
 @make_properties
 class SDFG(ControlFlowRegion):
     """ The main intermediate representation of code in DaCe.
@@ -458,6 +489,11 @@ class SDFG(ControlFlowRegion):
                        to_json=_arrays_to_json,
                        from_json=_nested_arrays_from_json)
     symbols = DictProperty(str, dtypes.typeclass, desc="Global symbols for this SDFG")
+    symbol_assumptions = DictProperty(str,
+                                      dict,
+                                      default={},
+                                      desc="Per-symbol SymPy assumption facts. Symbols STORED in the SDFG are "
+                                      "always bare -- assumptions live here and are applied transiently.")
 
     instrument = EnumProperty(dtype=dtypes.InstrumentationType,
                               desc="Measure execution statistics with given method",
@@ -520,6 +556,7 @@ def __init__(self,
         self._propagate = propagate
         self._parent = parent
         self.symbols = {}
+        self.symbol_assumptions = {}
         self._parent_sdfg = None
         self._parent_nsdfg_node = None
         self._arrays = NestedDict()  # type: Dict[str, dt.Array]
@@ -859,6 +896,7 @@ def replace_dict(self,
                 if validate_name(new_name):
                     _replace_dict_keys(self._arrays, name, new_name)
                     _replace_dict_keys(self.symbols, name, new_name)
+                    _replace_dict_keys(self.symbol_assumptions, name, new_name)
                     _replace_dict_keys(self.constants_prop, name, new_name)
                     _replace_dict_keys(self.callback_mapping, name, new_name)
                     _replace_dict_values(self.callback_mapping, name, new_name)
@@ -890,12 +928,59 @@ def add_symbol(self, name, stype, find_new_name: bool = False):
         self.symbols[name] = stype
         return name
 
+    def update_symbol_assumptions(self, name: str, **facts):
+        """Refine what is assumed about ``name``; the ONLY way to write the registry.
+
+        :param name: A symbol already declared with :meth:`add_symbol`.
+        :param facts: SymPy assumption facts, e.g. ``positive=True``. ``None`` values are ignored.
+        :return: ``None``. No graph object is touched; stored symbols stay bare.
+        :rtype: None
+        :raises KeyError: If ``name`` is not a declared symbol.
+        :raises ValueError: If a fact contradicts a recorded one; refinement is monotone.
+        """
+        if name not in self.symbols:
+            raise KeyError(f'Cannot assume anything about "{name}", it is not a symbol of this SDFG')
+        known = self.symbol_assumptions.setdefault(name, {})
+        for fact, value in facts.items():
+            if value is None:
+                continue
+            if known.get(fact, value) != value:
+                raise ValueError(f'Assumption "{fact}={value}" for symbol "{name}" contradicts the recorded '
+                                 f'"{fact}={known[fact]}"')
+            known[fact] = value
+
+    def assume_symbols(self, expr):
+        """``expr`` with bare symbols re-minted from the registry, for REASONING only.
+
+        This is the inverse of :func:`~dace.symbolic.bare_symbols` and the seam every site that
+        asks SymPy a question about a stored expression has to go through: a bare symbol answers
+        ``is_positive`` with ``None``, so ``Max``, ``sqrt`` and inequality folding all take their
+        conservative branch until the recorded facts are put back on.
+
+        :param expr: A symbolic expression, or any value, which is returned unchanged.
+        :return: ``expr`` with each registered symbol carrying its recorded assumptions, or
+                 ``expr`` unchanged when it is not symbolic. The assumed objects are local to
+                 the result and must never be stored back.
+        :rtype: Any
+        """
+        if not self.symbol_assumptions or not isinstance(expr, sympy.Basic):
+            return expr
+        repl = {}
+        for sym in expr.free_symbols:
+            facts = self.symbol_assumptions.get(sym.name)
+            if facts:
+                repl[sym] = symbolic.assumed_symbol(sym.name, self.symbols.get(sym.name), tuple(sorted(facts.items())))
+        # `xreplace` and not `subs`: the keys ARE the stored objects, so no matching is needed and
+        # none of SymPy's rewriting on the way out is wanted.
+        return expr.xreplace(repl) if repl else expr
+
     def remove_symbol(self, name):
         """ Removes a symbol from the SDFG.
 
             :param name: Symbol name.
         """
         del self.symbols[name]
+        self.symbol_assumptions.pop(name, None)
         # Clean up from symbol mapping if this SDFG is nested
         nsdfg = self.parent_nsdfg_node
         if nsdfg is not None and name in nsdfg.symbol_mapping:
@@ -2195,6 +2280,7 @@ def _add_symbols(sdfg: SDFG, desc: dt.Data):
         # Add the data descriptor to the SDFG and all symbols that are not yet known.
         self._arrays[name] = datadesc
         _add_symbols(self, datadesc)
+        absorb_symbol_assumptions(self, [datadesc])
 
         return name