Look for type information from other instructions in case of self type dependency

loopy/type_inference.py

@@ -699,6 +699,7 @@ def _infer_var_type(kernel, var_name, type_inf_mapper, subst_expander):
     import loopy as lp
 
     type_inf_mapper = type_inf_mapper.copy()
+    raised_type_dep_failure = False
 
     for writer_insn_id in kernel.writer_map().get(var_name, []):
         writer_insn = kernel.id_to_insn[writer_insn_id]
@@ -708,30 +709,38 @@ def _infer_var_type(kernel, var_name, type_inf_mapper, subst_expander):
         expr = subst_expander(writer_insn.expression)
 
         debug("             via expr %s", expr)
-        if isinstance(writer_insn, lp.Assignment):
-            result = type_inf_mapper(expr, return_dtype_set=True)
-        elif isinstance(writer_insn, lp.CallInstruction):
-            return_dtype_sets = type_inf_mapper(expr, return_tuple=True,
-                    return_dtype_set=True)
-
-            result = []
-            for return_dtype_set in return_dtype_sets:
-                result_i = None
-                found = False
-                for assignee, comp_dtype_set in zip(
-                        writer_insn.assignee_var_names(), return_dtype_set):
-                    if assignee == var_name:
-                        found = True
-                        result_i = comp_dtype_set
-                        break
 
-                assert found
-                if result_i is not None:
-                    result.append(result_i)
+        result = ()
+        try:
+            if isinstance(writer_insn, lp.Assignment):
+                result = type_inf_mapper(expr, return_dtype_set=True)
+            elif isinstance(writer_insn, lp.CallInstruction):
+                return_dtype_sets = type_inf_mapper(expr, return_tuple=True,
+                        return_dtype_set=True)
+
+                result = []
+                for return_dtype_set in return_dtype_sets:
+                    result_i = None
+                    found = False
+                    for assignee, comp_dtype_set in zip(
+                            writer_insn.assignee_var_names(), return_dtype_set):
+                        if assignee == var_name:
+                            found = True
+                            result_i = comp_dtype_set
+                            break
+
+                    assert found
+                    if result_i is not None:
+                        result.append(result_i)
+        except DependencyTypeInferenceFailure:
+            raised_type_dep_failure = True
 
-        debug("             result: %s", result)
+        if result:
+            debug("             result: %s", result)
+            dtype_sets.append(result)
 
-        dtype_sets.append(result)
+    if raised_type_dep_failure and not dtype_sets:
+        raise DependencyTypeInferenceFailure
 
     if not dtype_sets:
         return (

test/test_loopy.py

@@ -3216,6 +3216,20 @@ def test_get_return_from_kernel_mapping():
     assert ret_from_knl_idx[9] == 10
 
 
+def test_type_inf_with_self_deps(ctx_factory):
+    # See https://github.com/inducer/loopy/issues/481
+    ctx = ctx_factory()
+
+    prog = lp.make_kernel(
+            "{[i]: 0<=i<3}",
+            """
+            <> x[i] = if(i==0, 1, if(i==1, 1.1, if(i==2, 1.11, -1))) {id=insn1}
+            x[i] = if(x[i]>1.1, 2, x[i]) {id=insn2,dep=insn1,dup=i}
+            out[i] = x[i] {dep=insn2,dup=i}
+            """)
+    lp.auto_test_vs_ref(prog, ctx, prog)
+
+
 if __name__ == "__main__":
     if len(sys.argv) > 1:
         exec(sys.argv[1])