Change symbol typificaiton: Infer unknown LHS data types from RHS

src/pystencils/backend/kernelcreation/context.py

@@ -12,7 +12,7 @@ from ...sympyextensions.typed_sympy import TypedSymbol
 
 from ..symbols import PsSymbol
 from ..arrays import PsLinearizedArray
-from ...types import PsType, PsIntegerType, PsNumericType, PsScalarType, PsStructType
+from ...types import PsType, PsIntegerType, PsNumericType, PsScalarType, PsStructType, deconstify
 from ..constraints import KernelParamsConstraint
 from ..exceptions import PsInternalCompilerError, KernelConstraintsError
 
@@ -63,8 +63,8 @@ class KernelCreationContext:
         default_dtype: PsNumericType = DEFAULTS.numeric_dtype,
         index_dtype: PsIntegerType = DEFAULTS.index_dtype,
     ):
-        self._default_dtype = default_dtype
-        self._index_dtype = index_dtype
+        self._default_dtype = deconstify(default_dtype)
+        self._index_dtype = deconstify(index_dtype)
 
         self._symbols: dict[str, PsSymbol] = dict()
 

src/pystencils/backend/kernelcreation/typification.py

@@ -68,10 +68,18 @@ class TypeContext:
      - A set of restrictions on the target type:
        - `require_nonconst` to make sure the target type is not `const`, as required on assignment left-hand sides
        - Additional restrictions may be added in the future.
+    
+    The type context also tracks the tree traversal of the typifier:
+    
+     - ``is_lhs`` is set to True while a left-hand side expression is being processed,
+       and `False` while a right-hand side expression is processed.
     """
 
     def __init__(
-        self, target_type: PsType | None = None, require_nonconst: bool = False
+        self,
+        target_type: PsType | None = None,
+        require_nonconst: bool = False,
+        is_lhs: bool = False
     ):
         self._require_nonconst = require_nonconst
         self._deferred_exprs: list[PsExpression] = []
@@ -80,6 +88,8 @@ class TypeContext:
             self._fix_constness(target_type) if target_type is not None else None
         )
 
+        self._is_lhs = is_lhs
+
     @property
     def target_type(self) -> PsType | None:
         return self._target_type
@@ -88,6 +98,14 @@ class TypeContext:
     def require_nonconst(self) -> bool:
         return self._require_nonconst
     
+    @property
+    def is_lhs(self) -> bool:
+        return self._is_lhs
+    
+    @is_lhs.setter
+    def is_lhs(self, value: bool):
+        self._is_lhs = value
+
     def apply_dtype(self, dtype: PsType, expr: PsExpression | None = None):
         """Applies the given ``dtype`` to this type context, and optionally to the given expression.
 
@@ -171,8 +189,10 @@ class TypeContext:
                         )
 
                 case PsSymbolExpr(symb):
-                    assert symb.dtype is not None
-                    if not self._compatible(symb.dtype):
+                    if symb.dtype is None:
+                        #   Symbols are not forced to constness
+                        symb.dtype = deconstify(self._target_type)
+                    elif not self._compatible(symb.dtype):
                         raise TypificationError(
                             f"Type mismatch at symbol {symb}: Symbol type did not match the context's target type\n"
                             f"    Symbol type: {symb.dtype}\n"
@@ -262,7 +282,11 @@ class Typifier:
 
     The following general rules apply:
 
-     - The context's `default_dtype` is applied to all untyped symbols
+     - The context's ``default_dtype`` is applied to all untyped symbols encountered inside a right-hand side expression
+     - If an untyped symbol is encountered on an assignment's left-hand side, it will first be attempted to infer its
+       type from the right-hand side. If that fails, the context's ``default_dtype`` will be applied.
+     - It is an error if an untyped symbol occurs in the same type context as a typed symbol or constant
+       with a non-default data type.
      - By default, all expressions receive a ``const`` type unless they occur on a (non-declaration) assignment's
        left-hand side
 
@@ -280,7 +304,12 @@ class Typifier:
 
     def __call__(self, node: NodeT) -> NodeT:
         if isinstance(node, PsExpression):
-            self.visit_expr(node, TypeContext())
+            tc = TypeContext()
+            self.visit_expr(node, tc)
+
+            if tc.target_type is None:
+                #   no type could be inferred -> take the default
+                tc.apply_dtype(self._ctx.default_dtype)
         else:
             self.visit(node)
         return node
@@ -305,19 +334,31 @@ class Typifier:
 
             case PsDeclaration(lhs, rhs):
                 tc = TypeContext()
-                #   LHS defines target type; type context carries it to RHS
+
+                tc.is_lhs = True
                 self.visit_expr(lhs, tc)
-                assert tc.target_type is not None
+                tc.is_lhs = False
+
                 self.visit_expr(rhs, tc)
 
+                if tc.target_type is None:
+                    #   no type has been inferred -> use the default dtype
+                    tc.apply_dtype(self._ctx.default_dtype)
+
             case PsAssignment(lhs, rhs):
-                tc_lhs = TypeContext(require_nonconst=True)
+                tc_lhs = TypeContext(require_nonconst=True, is_lhs=True)
                 self.visit_expr(lhs, tc_lhs)
-                assert tc_lhs.target_type is not None
 
-                tc_rhs = TypeContext(tc_lhs.target_type, require_nonconst=False)
+                tc_rhs = TypeContext(target_type=tc_lhs.target_type)
                 self.visit_expr(rhs, tc_rhs)
 
+                if tc_rhs.target_type is None:
+                    tc_rhs.apply_dtype(self._ctx.default_dtype)
+
+                if tc_lhs.target_type is None:
+                    assert tc_rhs.target_type is not None
+                    tc_lhs.apply_dtype(deconstify(tc_rhs.target_type))
+
             case PsConditional(cond, branch_true, branch_false):
                 cond_tc = TypeContext(PsBoolType())
                 self.visit_expr(cond, cond_tc)
@@ -330,6 +371,7 @@ class Typifier:
             case PsLoop(ctr, start, stop, step, body):
                 if ctr.symbol.dtype is None:
                     ctr.symbol.apply_dtype(self._ctx.index_dtype)
+                    ctr.dtype = ctr.symbol.get_dtype()
 
                 tc_index = TypeContext(ctr.symbol.dtype)
                 self.visit_expr(start, tc_index)
@@ -355,11 +397,16 @@ class Typifier:
         either ``apply_dtype`` or ``infer_dtype``.
         """
         match expr:
-            case PsSymbolExpr(_):
-                if expr.symbol.dtype is None:
-                    expr.symbol.dtype = self._ctx.default_dtype
-
-                tc.apply_dtype(expr.symbol.dtype, expr)
+            case PsSymbolExpr(symb):
+                if tc.is_lhs:
+                    if symb.dtype is not None:
+                        tc.apply_dtype(symb.dtype, expr)
+                    elif tc.is_lhs:
+                        tc.infer_dtype(expr)
+                else:
+                    if symb.dtype is None:
+                        symb.dtype = self._ctx.default_dtype
+                    tc.apply_dtype(symb.dtype, expr)
 
             case PsConstantExpr(c):
                 if c.dtype is not None:

tests/nbackend/kernelcreation/test_typification.py

@@ -4,7 +4,7 @@ import numpy as np
 
 from typing import cast
 
-from pystencils import Assignment, TypedSymbol, Field, FieldType
+from pystencils import Assignment, TypedSymbol, Field, FieldType, AddAugmentedAssignment
 
 from pystencils.backend.ast.structural import (
     PsDeclaration,
@@ -186,7 +186,7 @@ def test_typify_structs():
         fasm = typify(fasm)
 
 
-def test_contextual_typing():
+def test_default_typing():
     ctx = KernelCreationContext()
     freeze = FreezeExpressions(ctx)
     typify = Typifier(ctx)
@@ -213,6 +213,43 @@ def test_contextual_typing():
     check(expr)
 
 
+def test_lhs_inference():
+    ctx = KernelCreationContext(default_dtype=create_numeric_type(np.float64))
+    freeze = FreezeExpressions(ctx)
+    typify = Typifier(ctx)
+
+    x, y, z = sp.symbols("x, y, z")
+    q = TypedSymbol("q", np.float32)
+    w = TypedSymbol("w", np.float16)
+
+    #   Type of the LHS is propagated to untyped RHS symbols
+
+    asm = Assignment(x, 3 - q)
+    fasm = typify(freeze(asm))
+
+    assert ctx.get_symbol("x").dtype == Fp(32)
+    assert fasm.lhs.dtype == constify(Fp(32))
+
+    asm = Assignment(y, 3 - w)
+    fasm = typify(freeze(asm))
+
+    assert ctx.get_symbol("y").dtype == Fp(16)
+    assert fasm.lhs.dtype == constify(Fp(16))
+
+    fasm = PsAssignment(PsExpression.make(ctx.get_symbol("z")), freeze(3 - w))
+    fasm = typify(fasm)
+
+    assert ctx.get_symbol("z").dtype == Fp(16)
+    assert fasm.lhs.dtype == Fp(16)
+
+    fasm = PsDeclaration(PsExpression.make(ctx.get_symbol("r")), PsLe(freeze(q), freeze(2 * q)))
+    fasm = typify(fasm)
+    
+    assert ctx.get_symbol("r").dtype == Bool()
+    assert fasm.lhs.dtype == constify(Bool())
+    assert fasm.rhs.dtype == constify(Bool())
+
+
 def test_erronous_typing():
     ctx = KernelCreationContext(default_dtype=create_numeric_type(np.float64))
     freeze = FreezeExpressions(ctx)
@@ -227,16 +264,23 @@ def test_erronous_typing():
     with pytest.raises(TypificationError):
         typify(expr)
 
+    #   Conflict between LHS and RHS symbols
     asm = Assignment(q, 3 - w)
     fasm = freeze(asm)
     with pytest.raises(TypificationError):
         typify(fasm)
 
+    #   Do not propagate types back from LHS symbols to RHS symbols
     asm = Assignment(q, 3 - x)
     fasm = freeze(asm)
     with pytest.raises(TypificationError):
         typify(fasm)
 
+    asm = AddAugmentedAssignment(z, 3 - q)
+    fasm = freeze(asm)
+    with pytest.raises(TypificationError):
+        typify(fasm)
+
 
 def test_typify_integer_binops():
     ctx = KernelCreationContext()

tests/symbolics/test_conditional_field_access.py

@@ -51,9 +51,6 @@ def add_fixed_constant_boundary_handling(assignments, with_cse):
 @pytest.mark.parametrize('dtype', ('float64', 'float32'))
 @pytest.mark.parametrize('with_cse', (False, 'with_cse'))
 def test_boundary_check(dtype, with_cse):
-    if with_cse:
-        pytest.xfail("Doesn't typify correctly yet.")
-
     f, g = ps.fields(f"f, g : {dtype}[2D]")
     stencil = ps.Assignment(g[0, 0], (f[1, 0] + f[-1, 0] + f[0, 1] + f[0, -1]) / 4)