From a55390595d26de1edbc11c0a3897f5e8a220e70e Mon Sep 17 00:00:00 2001 From: Andreas Kloeckner Date: Sat, 14 Aug 2021 15:16:45 -0500 Subject: [PATCH 1/2] Fix type inference for multiple assignments to a single variable (closes gh-481) Co-authored-by: Kaushik Kulkarni --- loopy/type_inference.py | 51 ++++++++++++++++++++++++----------------- 1 file changed, 30 insertions(+), 21 deletions(-) diff --git a/loopy/type_inference.py b/loopy/type_inference.py index f2fbe6d6f..351e9379c 100644 --- a/loopy/type_inference.py +++ b/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 ( From 81fc416d71b9385274662ae4de3b1d2876bb0160 Mon Sep 17 00:00:00 2001 From: Kaushik Kulkarni Date: Sat, 14 Aug 2021 15:18:52 -0500 Subject: [PATCH 2/2] test the type inference kernel with self deps Co-authored-by: Xiaoyu Wei --- test/test_loopy.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/test/test_loopy.py b/test/test_loopy.py index 7f7644f9a..63e8609f9 100644 --- a/test/test_loopy.py +++ b/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])