Repository: numba/numba
Pull request: Fix SSA to consider variables whose use is not dominated by the definition
Fixes #9242 and #9256.

* SSA reconstruction is adjusted to also consider definitions that do not dominate all uses.
* The above changes unfortunately is not compatible with Parfors; therefore, its effect is undone when Phi is stripped.
* https://github.com/numba/numba/pull/9255/commits/80dd8df240a4b73aa3af64b1abe077897dfa23a1 is part of debugging the Parfors breakage due to the initial commit. Any user code that makes a new variable definition in a loop not in the pattern of reduction will trigger a problem in SROA-avoidance due to incompatible postcondition of use-def analysis on Parfor node.

 

Problem statement:
Broken SSA for conditionally defined var
```
@njit
def py_func(a):
    c = a > 0
    if c:
        d = a + 5

    return c and d > 0

py_func(10)
```
`NUMBA_DUMP_SSA=1` gives
```
label 0:
    a = arg(0, name=a)                       ['a']
    $const4.1 = const(int, 0)                ['$const4.1']
    c = a > $const4.1                        ['$const4.1', 'a', 'c']
    bool12 = global(bool: <class 'bool'>)    ['bool12']
    $12pred = call bool12(c, func=bool12, args=(Var(c, test_basic.py:929),), kws=(), vararg=None, varkwarg=None, target=None) ['$12pred', 'bool12', 'c']
    branch $12pred, 14, 22                   ['$12pred']
label 14:
    $const16.1 = const(int, 5)               ['$const16.1']
    d = a + $const16.1                       ['$const16.1', 'a', 'd']
    jump 22                                  []
label 22:
    bool24 = global(bool: <class 'bool'>)    ['bool24']
    $24pred = call bool24(c, func=bool24, args=(Var(c, test_basic.py:929),), kws=(), vararg=None, varkwarg=None, target=None) ['$24pred', 'bool24', 'c']
    $phi32.0 = c                             ['$phi32.0', 'c']
    branch $24pred, 26, 32                   ['$24pred']
label 26:
    $const28.1 = const(int, 0)               ['$const28.1']
    $30compare_op.2 = d > $const28.1         ['$30compare_op.2', '$const28.1', 'd']
    $phi32.0.1 = $30compare_op.2             ['$30compare_op.2', '$phi32.0.1']
    jump 32                                  []
label 32:
    $phi32.0.2 = phi(incoming_values=[Var($phi32.0.1, test_basic.py:933), Var($phi32.0, test_basic.py:933)], incoming_blocks=[26, 22]) ['$phi32.0', '$phi32.0.1', '$phi32.0.2']
    $32return_value.1 = cast(value=$phi32.0.2) ['$32return_value.1', '$phi32.0.2']
    return $32return_value.1                 ['$32return_value.1']
```
`d` is used in block 26 and defined in block 14, which doesn't dominate 26

Hints from issue discussion:
As discussed with OP out-of-band, Numba's SSA pass is only adjusting variables with multiple definition. Here, it fails to recognize that `d` has an uninitialized path that should require a phi node.  The failure to recognize this affects subsequent compiler analysis relying on the SSA property.

Commit message: Apply suggestions from code review

Co-authored-by: stuartarchibald <stuartarchibald@users.noreply.github.com>

Patch to review:
diff --git a/docs/upcoming_changes/9242.bug_fix.rst b/docs/upcoming_changes/9242.bug_fix.rst
new file mode 100644
index 00000000000..1062af2155a
--- /dev/null
+++ b/docs/upcoming_changes/9242.bug_fix.rst
@@ -0,0 +1,6 @@
+Fix SSA to consider variables where use is not dominated by the definition
+==========================================================================
+
+A SSA problem is fixed such that a conditionally defined variable will receive
+a phi node showing that there is a path where the variable is undefined.
+This affects extension code that relies on SSA behavior.
diff --git a/numba/core/ssa.py b/numba/core/ssa.py
index 8a368229d99..a14c94585a2 100644
--- a/numba/core/ssa.py
+++ b/numba/core/ssa.py
@@ -54,7 +54,7 @@ def _run_ssa(blocks):
     cfg = compute_cfg_from_blocks(blocks)
     df_plus = _iterated_domfronts(cfg)
     # Find SSA violators
-    violators = _find_defs_violators(blocks)
+    violators = _find_defs_violators(blocks, cfg)
     # Make cache for .list_vars()
     cache_list_vars = _CacheListVars()
 
@@ -142,7 +142,7 @@ def _get_scope(blocks):
     return first.scope
 
 
-def _find_defs_violators(blocks):
+def _find_defs_violators(blocks, cfg):
     """
     Returns
     -------
@@ -150,9 +150,22 @@ def _find_defs_violators(blocks):
         The SSA violators in a dictionary of variable names.
     """
     defs = defaultdict(list)
-    _run_block_analysis(blocks, defs, _GatherDefsHandler())
+    uses = defaultdict(set)
+    states = dict(defs=defs, uses=uses)
+    _run_block_analysis(blocks, states, _GatherDefsHandler())
     _logger.debug("defs %s", pformat(defs))
+    # Gather violators by number of definitions
     violators = {k for k, vs in defs.items() if len(vs) > 1}
+    # Gather violators by uses not dominated by the one def
+    doms = cfg.dominators()
+    for k, use_blocks in uses.items():
+        if k not in violators:
+            for label in use_blocks:
+                dom = doms[label]
+                def_labels = {label for _assign, label in defs[k] }
+                if not def_labels.intersection(dom):
+                    violators.add(k)
+                    break
     _logger.debug("SSA violators %s", pformat(violators))
     return violators
 
@@ -160,6 +173,7 @@ def _find_defs_violators(blocks):
 def _run_block_analysis(blocks, states, handler):
     for label, blk in blocks.items():
         _logger.debug("==== SSA block analysis pass on %s", label)
+        states['label'] = label
         for _ in _run_ssa_block_pass(states, blk, handler):
             pass
 
@@ -241,12 +255,26 @@ def on_other(self, states, stmt):
 
 
 class _GatherDefsHandler(_BaseHandler):
-    """Find all defs
+    """Find all defs and uses of variable in each block
 
-    ``states`` is a Mapping[str, List[ir.Assign]]
+    ``states["label"]`` is a int; label of the current block
+    ``states["defs"]`` is a Mapping[str, List[Tuple[ir.Assign, int]]]
+    ``states["uses"]`` is a Mapping[Set[int]]
     """
     def on_assign(self, states, assign):
-        states[assign.target.name].append(assign)
+        # keep track of assignment and the block
+        states["defs"][assign.target.name].append((assign, states["label"]))
+        # keep track of uses
+        for var in assign.list_vars():
+            k = var.name
+            if k != assign.target.name:
+                states["uses"][k].add(states["label"])
+
+    def on_other(self, states, stmt):
+        # keep track of uses
+        for var in stmt.list_vars():
+            k = var.name
+            states["uses"][k].add(states["label"])
 
 
 class UndefinedVariable:
diff --git a/numba/core/typed_passes.py b/numba/core/typed_passes.py
index 061295cb4fb..a25dbc9c839 100644
--- a/numba/core/typed_passes.py
+++ b/numba/core/typed_passes.py
@@ -1,6 +1,7 @@
 import abc
 from contextlib import contextmanager
 from collections import defaultdict, namedtuple
+from functools import partial
 from copy import copy
 import warnings
 
@@ -23,7 +24,8 @@
                                  dead_code_elimination, simplify_CFG,
                                  get_definition,
                                  build_definitions, compute_cfg_from_blocks,
-                                 is_operator_or_getitem)
+                                 is_operator_or_getitem,
+                                 replace_vars)
 from numba.core import postproc
 from llvmlite import binding as llvm
 
@@ -881,6 +883,10 @@ def __init__(self):
     def run_pass(self, state):
         state.func_ir = self._strip_phi_nodes(state.func_ir)
         state.func_ir._definitions = build_definitions(state.func_ir.blocks)
+        if "flags" in state and state.flags.auto_parallel.enabled:
+            self._simplify_conditionally_defined_variable(state.func_ir)
+            state.func_ir._definitions = build_definitions(state.func_ir.blocks)
+
         # Rerun postprocessor to update metadata
         post_proc = postproc.PostProcessor(state.func_ir)
         post_proc.run(emit_dels=False)
@@ -956,3 +962,78 @@ def _strip_phi_nodes(self, func_ir):
 
         func_ir.blocks = newblocks
         return func_ir
+
+    def _simplify_conditionally_defined_variable(self, func_ir):
+        """
+        Rewrite assignments like:
+
+            ver1 = null()
+            ...
+            ver1 = ver
+            ...
+            uses(ver1)
+
+        into:
+            # delete all assignments to ver1
+            uses(ver)
+
+        This is only needed for parfors because the SSA pass will create extra
+        variable assignments that the parfor code does not expect.
+        This pass helps avoid problems by reverting the effect of SSA.
+        """
+        any_block = next(iter(func_ir.blocks.values()))
+        scope = any_block.scope
+        defs = func_ir._definitions
+
+        def unver_or_undef(unver, defn):
+            # Is the definition undefined or pointing to the unversioned name?
+            if isinstance(defn, ir.Var):
+                if defn.unversioned_name == unver:
+                    return True
+            elif isinstance(defn, ir.Expr):
+                if defn.op == "null":
+                    return True
+            return False
+
+        def legalize_all_versioned_names(var):
+            # Is all versioned names undefined or defined to the same
+            # variable chain?
+            if not var.versioned_names:
+                return False
+            for verioned in var.versioned_names:
+                vs = defs.get(verioned, ())
+                if not all(map(partial(unver_or_undef, k), vs)):
+                    return False
+            return True
+
+        # Find unversioned variables that met the conditions
+        suspects = set()
+        for k in defs:
+            try:
+                # This may fail?
+                var = scope.get_exact(k)
+            except errors.NotDefinedError:
+                continue
+            # is the unversioned?
+            if var.unversioned_name == k:
+                if legalize_all_versioned_names(var):
+                    suspects.add(var)
+
+        delete_set = set()
+        replace_map = {}
+        for var in suspects:
+            # rewrite Var uses to the unversioned name
+            for versioned in var.versioned_names:
+                ver_var = scope.get_exact(versioned)
+                # delete assignment to the versioned name
+                delete_set.add(ver_var)
+                # replace references to versioned name with the unversioned
+                replace_map[versioned] = var
+
+        # remove assignments to the versioned names
+        for _label, blk in func_ir.blocks.items():
+            for assign in blk.find_insts(ir.Assign):
+                if assign.target in delete_set:
+                    blk.remove(assign)
+        # do variable replacement
+        replace_vars(func_ir.blocks, replace_map)
diff --git a/numba/parfors/parfor_lowering.py b/numba/parfors/parfor_lowering.py
index 98aaa4b8362..f412ca40d1d 100644
--- a/numba/parfors/parfor_lowering.py
+++ b/numba/parfors/parfor_lowering.py
@@ -52,6 +52,15 @@ def lower_inst(self, inst):
         else:
             super().lower_inst(inst)
 
+    @property
+    def _disable_sroa_like_opt(self):
+        """
+        Force disable this because Parfor use-defs is incompatible---it only
+        considers use-defs in blocks that must be executing.
+        See https://github.com/numba/numba/commit/017e2ff9db87fc34149b49dd5367ecbf0bb45268
+        """
+        return True
+
 
 def _lower_parfor_parallel(lowerer, parfor):
     if parfor.lowerer is None:
diff --git a/numba/tests/test_parfors.py b/numba/tests/test_parfors.py
index 2fc5db6feab..9ad8ab7c9dc 100644
--- a/numba/tests/test_parfors.py
+++ b/numba/tests/test_parfors.py
@@ -2311,6 +2311,44 @@ def test_impl(a, b):
         self.check(test_impl, 3.7, 4.3)
         self.assertEqual(countParfors(test_impl, (types.float64, types.float64)), 1)
 
+    def test_issue9256_lower_sroa_conflict(self):
+        @njit(parallel=True)
+        def def_in_loop(x):
+            c = 0
+            set_num_threads(1)
+            for i in prange(x):
+                c = i
+            return c
+
+        self.assertEqual(def_in_loop(10), def_in_loop.py_func(10))
+
+    def test_issue9256_lower_sroa_conflict_variant1(self):
+        def def_in_loop(x):
+            c = x
+            set_num_threads(1)
+            for _i in prange(x):
+                if c: # forces 3 SSA versions
+                    d = x + 4
+            return c, d > 0
+
+        expected = def_in_loop(4)
+        self.assertEqual(expected, njit(parallel=False)(def_in_loop)(4))
+        self.assertEqual(expected, njit(parallel=True)(def_in_loop)(4))
+
+    def test_issue9256_lower_sroa_conflict_variant2(self):
+        def def_in_loop(x):
+            c = x
+            set_num_threads(1)
+            for _i in prange(x):
+                if c:
+                    for _j in range(x): # forces 4 SSA versions
+                        d = x + 4
+            return c, d > 0
+
+        expected = def_in_loop(4)
+        self.assertEqual(expected, njit(parallel=False)(def_in_loop)(4))
+        self.assertEqual(expected, njit(parallel=True)(def_in_loop)(4))
+
 
 @skip_parfors_unsupported
 class TestParforsLeaks(MemoryLeakMixin, TestParforsBase):
diff --git a/numba/tests/test_ssa.py b/numba/tests/test_ssa.py
index 61642d960e9..49e28c05f2a 100644
--- a/numba/tests/test_ssa.py
+++ b/numba/tests/test_ssa.py
@@ -498,6 +498,60 @@ def while_for(n, max_iter=1):
         # One phi?
         self.assertEqual(phi_counter, [1])
 
+    def test_issue9242_use_not_dom_def(self):
+        from numba.core.ir import FunctionIR
+        from numba.core.compiler_machinery import (
+            AnalysisPass,
+            register_pass,
+        )
+
+        def check(fir: FunctionIR):
+            [blk, *_] = fir.blocks.values()
+            var = blk.scope.get("d")
+            defn = fir.get_definition(var)
+            self.assertEqual(defn.op, "phi")
+            self.assertIn(ir.UNDEFINED, defn.incoming_values)
+
+        @register_pass(mutates_CFG=False, analysis_only=True)
+        class SSACheck(AnalysisPass):
+            """
+            Check SSA on variable `d`
+            """
+
+            _name = "SSA_Check"
+
+            def __init__(self):
+                AnalysisPass.__init__(self)
+
+            def run_pass(self, state):
+                check(state.func_ir)
+                return False
+
+        class SSACheckPipeline(CompilerBase):
+            """Inject SSACheck pass into the default pipeline following the SSA
+            pass
+            """
+
+            def define_pipelines(self):
+                pipeline = DefaultPassBuilder.define_nopython_pipeline(
+                    self.state, "ssa_check_custom_pipeline")
+
+                pipeline._finalized = False
+                pipeline.add_pass_after(SSACheck, ReconstructSSA)
+
+                pipeline.finalize()
+                return [pipeline]
+
+        @njit(pipeline_class=SSACheckPipeline)
+        def py_func(a):
+            c = a > 0
+            if c:
+                d = a + 5  # d is only defined here; undef in the else branch
+
+            return c and d > 0
+
+        py_func(10)
+
 
 class TestSROAIssues(MemoryLeakMixin, TestCase):
     # This tests issues related to the SROA optimization done in lowering, which


Code review comments:
diff_hunk: @@ -956,3 +962,78 @@ def _strip_phi_nodes(self, func_ir):
 
         func_ir.blocks = newblocks
         return func_ir
+
+    def _simplify_conditionally_defined_variable(self, func_ir):
+        """
+        Rewrite assignments like:
+
+            ver1 = null()
+            ...
+            ver1 = ver
+            ...
+            uses(ver1)
+
+        into:
+            # delete all assignments to ver1
+            uses(ver)
+
+        This is only needed for parfors because the SSA pass will create extra
+        variable assignments that the parfor code does not expect.
+        This pass helps avoid problems by reverting the effect of SSA.
+        """
+        any_block = next(iter(func_ir.blocks.values()))
+        scope = any_block.scope
+        defs = func_ir._definitions
+
+        def unver_or_undef(unver, defn):
+            # Is the definition undefined or pointing to the unversioned name?
+            if isinstance(defn, ir.Var):
+                if defn.unversioned_name == unver:
+                    return True
+            elif isinstance(defn, ir.Expr):
+                if defn.op == "null":
+                    return True
+            return False
+
+        def legalize_all_versioned_names(var):
+            # Is all versioned names undefined or defined to the same
+            # variable chain?
+            if not var.versioned_names:
+                return False
+            for verioned in var.versioned_names:
+                vs = defs.get(verioned, ())
+                if not all(map(partial(unver_or_undef, k), vs)):
+                    return False
+            return True
+
+        # Find unversioned variables that met the conditions
+        suspects = set()
+        for k in defs:
+            try:
+                # This may fail?
+                var = scope.get_exact(k)
+            except errors.NotDefinedError:
+                continue
+            # is the unversioned?
path: numba/core/typed_passes.py
text: @user1:
```suggestion
            # is the var name unversioned?
```

diff_hunk: @@ -956,3 +962,78 @@ def _strip_phi_nodes(self, func_ir):
 
         func_ir.blocks = newblocks
         return func_ir
+
+    def _simplify_conditionally_defined_variable(self, func_ir):
+        """
+        Rewrite assignments like:
+
+            ver1 = null()
+            ...
+            ver1 = ver
+            ...
+            uses(ver1)
+
+        into:
+            # delete all assignments to ver1
+            uses(ver)
+
+        This is only needed for parfors because the SSA pass will create extra
+        variable assignments that the parfor code does not expect.
+        This pass helps avoid problems by reverting the effect of SSA.
+        """
+        any_block = next(iter(func_ir.blocks.values()))
+        scope = any_block.scope
+        defs = func_ir._definitions
+
+        def unver_or_undef(unver, defn):
+            # Is the definition undefined or pointing to the unversioned name?
+            if isinstance(defn, ir.Var):
+                if defn.unversioned_name == unver:
+                    return True
+            elif isinstance(defn, ir.Expr):
+                if defn.op == "null":
+                    return True
+            return False
+
+        def legalize_all_versioned_names(var):
+            # Is all versioned names undefined or defined to the same
+            # variable chain?
+            if not var.versioned_names:
+                return False
+            for verioned in var.versioned_names:
+                vs = defs.get(verioned, ())
path: numba/core/typed_passes.py
text: @user1:
```suggestion
            for versioned in var.versioned_names:
                vs = defs.get(versioned, ())
```

diff_hunk: @@ -956,3 +962,78 @@ def _strip_phi_nodes(self, func_ir):
 
         func_ir.blocks = newblocks
         return func_ir
+
+    def _simplify_conditionally_defined_variable(self, func_ir):
+        """
+        Rewrite assignments like:
+
+            ver1 = null()
+            ...
+            ver1 = ver
+            ...
+            uses(ver1)
+
+        into:
+            # delete all assignments to ver1
+            uses(ver)
+
+        This is only needed for parfors because the SSA pass will create extra
+        variable assignments that the parfor code does not expect.
+        This pass helps avoid problems by reverting the effect of SSA.
+        """
+        any_block = next(iter(func_ir.blocks.values()))
+        scope = any_block.scope
+        defs = func_ir._definitions
+
+        def unver_or_undef(unver, defn):
+            # Is the definition undefined or pointing to the unversioned name?
+            if isinstance(defn, ir.Var):
+                if defn.unversioned_name == unver:
+                    return True
+            elif isinstance(defn, ir.Expr):
+                if defn.op == "null":
+                    return True
+            return False
+
+        def legalize_all_versioned_names(var):
+            # Is all versioned names undefined or defined to the same
path: numba/core/typed_passes.py
text: @user1:
```suggestion
            # Are all versioned names undefined or defined to the same
```