From 9290fd20b409ce0f865351280a44ba215e3a0eba Mon Sep 17 00:00:00 2001 From: Lukas Truemper Date: Tue, 6 Jun 2023 21:06:53 +0200 Subject: [PATCH 1/3] Bugfix in taskletfusion --- .../transformation/dataflow/tasklet_fusion.py | 3 +++ tests/transformations/tasklet_fusion_test.py | 23 +++++++++++++++++++ 2 files changed, 26 insertions(+) diff --git a/dace/transformation/dataflow/tasklet_fusion.py b/dace/transformation/dataflow/tasklet_fusion.py index 8179ead457..398af7c8b9 100644 --- a/dace/transformation/dataflow/tasklet_fusion.py +++ b/dace/transformation/dataflow/tasklet_fusion.py @@ -199,6 +199,9 @@ def apply(self, graph: dace.SDFGState, sdfg: dace.SDFG): repldict = {} for in_edge in graph.in_edges(t1): old_value = in_edge.dst_conn + if old_value is None: + continue + # Check if there is a conflict. if in_edge.dst_conn in inputs: # Conflicts are ok if the Memlets are the same. diff --git a/tests/transformations/tasklet_fusion_test.py b/tests/transformations/tasklet_fusion_test.py index a65d218d98..8c5e06ed58 100644 --- a/tests/transformations/tasklet_fusion_test.py +++ b/tests/transformations/tasklet_fusion_test.py @@ -178,6 +178,28 @@ def test_tasklet_fusion_multiline(A: datatype): assert (result[0] == 11) +def test_map_param(): + @dace.program + def map_uses_param(A: dace.float32[10], B: dace.float32[10], C: dace.float32[10]): + for i in dace.map[0:10]: + a = i - A[i] + b = B[i] * i + C[i] = a + b + + sdfg = map_uses_param.to_sdfg(simplify=True) + + num_tasklet_fusions = sdfg.apply_transformations(TaskletFusion) + assert (num_tasklet_fusions == 1) + + A = np.zeros([10], dtype=np.float32) + B = np.ones([10], dtype=np.float32) + C = np.empty([10], dtype=np.float32) + sdfg(A=A, B=B, C=C) + + ref = np.array(range(0, 10, 1)) * 2.0 + assert (C == ref).all() + + @pytest.mark.parametrize('with_data', [pytest.param(True), pytest.param(False)]) @pytest.mark.parametrize('language', [pytest.param('CPP'), pytest.param('Python')]) def test_map_with_tasklets(language: str, with_data: bool): @@ -200,6 +222,7 @@ def test_map_with_tasklets(language: str, with_data: bool): test_same_name() test_same_name_different_memlet() test_tasklet_fusion_multiline() + test_map_param() test_map_with_tasklets(language='Python', with_data=False) test_map_with_tasklets(language='Python', with_data=True) test_map_with_tasklets(language='CPP', with_data=False) From c3f5548aacdc0ccec921a9657a70cf85ed1e95f1 Mon Sep 17 00:00:00 2001 From: Lukas Truemper Date: Tue, 6 Jun 2023 21:07:31 +0200 Subject: [PATCH 2/3] yapf in taskletfusion classes --- .../transformation/dataflow/tasklet_fusion.py | 25 +++++----------- tests/transformations/tasklet_fusion_test.py | 29 +++++++++---------- 2 files changed, 20 insertions(+), 34 deletions(-) diff --git a/dace/transformation/dataflow/tasklet_fusion.py b/dace/transformation/dataflow/tasklet_fusion.py index 398af7c8b9..99f8f625be 100644 --- a/dace/transformation/dataflow/tasklet_fusion.py +++ b/dace/transformation/dataflow/tasklet_fusion.py @@ -32,7 +32,6 @@ def visit_Name(self, node: ast.Name) -> Any: class CPPConnectorRenamer(): - def __init__(self, repl_dict: Dict[str, str]) -> None: self.repl_dict = repl_dict @@ -44,7 +43,6 @@ def rename(self, code: str) -> str: class PythonInliner(ast.NodeTransformer): - def __init__(self, target_id, target_ast): self.target_id = target_id self.target_ast = target_ast @@ -57,7 +55,6 @@ def visit_Name(self, node: ast.AST): class CPPInliner(): - def __init__(self, inline_target, inline_val): self.inline_target = inline_target self.inline_val = inline_val @@ -144,10 +141,7 @@ class TaskletFusion(pm.SingleStateTransformation): @classmethod def expressions(cls): - return [ - sdutil.node_path_graph(cls.t1, cls.data, cls.t2), - sdutil.node_path_graph(cls.t1, cls.t2) - ] + return [sdutil.node_path_graph(cls.t1, cls.data, cls.t2), sdutil.node_path_graph(cls.t1, cls.t2)] def can_be_applied(self, graph: dace.SDFGState, expr_index: int, sdfg: dace.SDFG, permissive: bool = False) -> bool: t1 = self.t1 @@ -191,9 +185,7 @@ def apply(self, graph: dace.SDFGState, sdfg: dace.SDFG): t2_in_edge = graph.out_edges(data if data is not None else t1)[0] # Remove the connector from the second Tasklet. - inputs = { - k: v for k, v in t2.in_connectors.items() if k != t2_in_edge.dst_conn - } + inputs = {k: v for k, v in t2.in_connectors.items() if k != t2_in_edge.dst_conn} # Copy the first Tasklet's in connectors. repldict = {} @@ -214,8 +206,8 @@ def apply(self, graph: dace.SDFGState, sdfg: dace.SDFG): break else: t2edge = conflict_edges[0] - if t2edge is not None and (in_edge.data != t2edge.data or in_edge.data.data != t2edge.data.data or - in_edge.data is None or in_edge.data.data is None): + if t2edge is not None and (in_edge.data != t2edge.data or in_edge.data.data != t2edge.data.data + or in_edge.data is None or in_edge.data.data is None): in_edge.dst_conn = dace.data.find_new_name(in_edge.dst_conn, set(inputs)) repldict[old_value] = in_edge.dst_conn else: @@ -231,9 +223,7 @@ def apply(self, graph: dace.SDFGState, sdfg: dace.SDFG): if repldict: assigned_value = PythonConnectorRenamer(repldict).visit(assigned_value) - new_code = [ - PythonInliner(t2_in_edge.dst_conn, assigned_value).visit(line) for line in t2.code.code - ] + new_code = [PythonInliner(t2_in_edge.dst_conn, assigned_value).visit(line) for line in t2.code.code] new_code_str = '\n'.join(astunparse.unparse(line) for line in new_code) elif t1.language == Language.CPP: assigned_value = t1.code.as_string @@ -255,9 +245,8 @@ def apply(self, graph: dace.SDFGState, sdfg: dace.SDFG): else: return - new_tasklet = graph.add_tasklet( - t1.label + '_fused_' + t2.label, inputs, t2.out_connectors, new_code_str, t1.language - ) + new_tasklet = graph.add_tasklet(t1.label + '_fused_' + t2.label, inputs, t2.out_connectors, new_code_str, + t1.language) for in_edge in graph.in_edges(t1): graph.add_edge(in_edge.src, in_edge.src_conn, new_tasklet, in_edge.dst_conn, in_edge.data) diff --git a/tests/transformations/tasklet_fusion_test.py b/tests/transformations/tasklet_fusion_test.py index 8c5e06ed58..1e10759753 100644 --- a/tests/transformations/tasklet_fusion_test.py +++ b/tests/transformations/tasklet_fusion_test.py @@ -10,6 +10,7 @@ M = 10 N = 2 * M + @dace.program def map_with_tasklets(A: datatype[N], B: datatype[M]): C = np.zeros_like(B) @@ -42,15 +43,11 @@ def _make_sdfg(language: str, with_data: bool = False): outputs = { '__out': datatype, } - ta = state.add_tasklet( - 'a', inputs, { - '__out1': datatype, - '__out2': datatype, - '__out3': datatype, - }, - f'__out1 = __inp1 + __inp2{endl}__out2 = __out1{endl}__out3 = __out1{endl}', - lang - ) + ta = state.add_tasklet('a', inputs, { + '__out1': datatype, + '__out2': datatype, + '__out3': datatype, + }, f'__out1 = __inp1 + __inp2{endl}__out2 = __out1{endl}__out3 = __out1{endl}', lang) tb = state.add_tasklet('b', inputs, outputs, f'__out = __inp1 * __inp2{endl}', lang) tc = state.add_tasklet('c', inputs, outputs, f'__out = __inp1 + __inp2{endl}', lang) td = state.add_tasklet('d', inputs, outputs, f'__out = __inp1 / __inp2{endl}', lang) @@ -60,12 +57,12 @@ def _make_sdfg(language: str, with_data: bool = False): state.add_memlet_path(A, me, tb, memlet=dace.Memlet('A[2*i]'), dst_conn='__inp2') state.add_memlet_path(B, me, tc, memlet=dace.Memlet('B[i]'), dst_conn='__inp2') if with_data: - sdfg.add_array('tmp1', (1,), datatype, dtypes.StorageType.Default, None, True) - sdfg.add_array('tmp2', (1,), datatype, dtypes.StorageType.Default, None, True) - sdfg.add_array('tmp3', (1,), datatype, dtypes.StorageType.Default, None, True) - sdfg.add_array('tmp4', (1,), datatype, dtypes.StorageType.Default, None, True) - sdfg.add_array('tmp5', (1,), datatype, dtypes.StorageType.Default, None, True) - sdfg.add_array('tmp6', (1,), datatype, dtypes.StorageType.Default, None, True) + sdfg.add_array('tmp1', (1, ), datatype, dtypes.StorageType.Default, None, True) + sdfg.add_array('tmp2', (1, ), datatype, dtypes.StorageType.Default, None, True) + sdfg.add_array('tmp3', (1, ), datatype, dtypes.StorageType.Default, None, True) + sdfg.add_array('tmp4', (1, ), datatype, dtypes.StorageType.Default, None, True) + sdfg.add_array('tmp5', (1, ), datatype, dtypes.StorageType.Default, None, True) + sdfg.add_array('tmp6', (1, ), datatype, dtypes.StorageType.Default, None, True) atemp1 = state.add_access('tmp1') atemp2 = state.add_access('tmp2') atemp3 = state.add_access('tmp3') @@ -101,7 +98,7 @@ def test_basic(): def test_basic_tf(A: datatype[5, 5]): B = A + 1 return B * 2 - + sdfg = test_basic_tf.to_sdfg(simplify=True) num_map_fusions = sdfg.apply_transformations(MapFusion) From 642bf2a2a4be9f3273af1306990fb6f4cca346de Mon Sep 17 00:00:00 2001 From: Lukas Truemper Date: Tue, 6 Jun 2023 21:11:21 +0200 Subject: [PATCH 3/3] Corrected taskletfusion test case --- tests/transformations/tasklet_fusion_test.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/transformations/tasklet_fusion_test.py b/tests/transformations/tasklet_fusion_test.py index 1e10759753..c7fd6802d5 100644 --- a/tests/transformations/tasklet_fusion_test.py +++ b/tests/transformations/tasklet_fusion_test.py @@ -185,8 +185,8 @@ def map_uses_param(A: dace.float32[10], B: dace.float32[10], C: dace.float32[10] sdfg = map_uses_param.to_sdfg(simplify=True) - num_tasklet_fusions = sdfg.apply_transformations(TaskletFusion) - assert (num_tasklet_fusions == 1) + num_tasklet_fusions = sdfg.apply_transformations_repeated(TaskletFusion) + assert (num_tasklet_fusions == 3) A = np.zeros([10], dtype=np.float32) B = np.ones([10], dtype=np.float32)