From 8ff0cc7acbe4df2b15e30010f98a483790b11780 Mon Sep 17 00:00:00 2001 From: Ahmad Nawab Date: Mon, 2 Oct 2023 16:33:42 +0200 Subject: [PATCH 1/2] Add tests to check correct placement of loop-bound assignment --- .../tests/test_single_column_coalesced.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/transformations/tests/test_single_column_coalesced.py b/transformations/tests/test_single_column_coalesced.py index 619a38f52..62c4b4d26 100644 --- a/transformations/tests/test_single_column_coalesced.py +++ b/transformations/tests/test_single_column_coalesced.py @@ -567,6 +567,13 @@ def test_scc_hoist_multiple_kernels_loops(frontend, horizontal, vertical, blocki assert driver_loops[2].variable == 'jk' assert driver_loops[2].bounds == '2:nz' + # check location of loop-bound assignment + assign = FindNodes(Assignment).visit(driver_loops[0])[0] + assert assign.lhs == 'end' + assert assign.rhs == 'nlon-nb' + assigns = FindNodes(Assignment).visit(driver_loops[1]) + assert not assign in assigns + assert driver_loops[4] in FindNodes(Loop).visit(driver_loops[3].body) assert driver_loops[5] in FindNodes(Loop).visit(driver_loops[3].body) assert driver_loops[6] in FindNodes(Loop).visit(driver_loops[3].body) @@ -596,6 +603,13 @@ def test_scc_hoist_multiple_kernels_loops(frontend, horizontal, vertical, blocki assert driver_loops[10].variable == 'jk' assert driver_loops[10].bounds == '2:nz' + # check location of loop-bound assignment + assign = FindNodes(Assignment).visit(driver_loops[8])[0] + assert assign.lhs == 'end' + assert assign.rhs == 'nlon-nb' + assigns = FindNodes(Assignment).visit(driver_loops[9]) + assert not assign in assigns + rmtree(basedir) From 6e615f07ef2da766f057de09cc850784542819bf Mon Sep 17 00:00:00 2001 From: Ahmad Nawab Date: Mon, 2 Oct 2023 16:43:37 +0200 Subject: [PATCH 2/2] Extend trim_vector_sections functionality to driver loop --- transformations/tests/test_single_column_coalesced.py | 11 +++++++---- .../transformations/single_column_coalesced_vector.py | 2 +- 2 files changed, 8 insertions(+), 5 deletions(-) diff --git a/transformations/tests/test_single_column_coalesced.py b/transformations/tests/test_single_column_coalesced.py index 62c4b4d26..8294adcb6 100644 --- a/transformations/tests/test_single_column_coalesced.py +++ b/transformations/tests/test_single_column_coalesced.py @@ -410,7 +410,8 @@ def test_scc_hoist_multiple_kernels(frontend, horizontal, vertical, blocking): @pytest.mark.parametrize('frontend', available_frontends()) -def test_scc_hoist_multiple_kernels_loops(frontend, horizontal, vertical, blocking): +@pytest.mark.parametrize('trim_vector_sections', [True, False]) +def test_scc_hoist_multiple_kernels_loops(frontend, trim_vector_sections, horizontal, vertical, blocking): """ Test hoisting of column temporaries to "driver" level. """ @@ -429,7 +430,6 @@ def test_scc_hoist_multiple_kernels_loops(frontend, horizontal, vertical, blocki !$loki driver-loop do b=1, nb end = nlon - nb - !$loki separator do jk = 2, nz do jl = start, end q(jl, jk, b) = 2.0 * jk * jl @@ -523,7 +523,7 @@ def test_scc_hoist_multiple_kernels_loops(frontend, horizontal, vertical, blocki kernel = scheduler.item_map["kernel_mod#kernel"].routine transformation = (SCCBaseTransformation(horizontal=horizontal, directive='openacc'),) - transformation += (SCCDevectorTransformation(horizontal=horizontal, trim_vector_sections=False),) + transformation += (SCCDevectorTransformation(horizontal=horizontal, trim_vector_sections=trim_vector_sections),) transformation += (SCCDemoteTransformation(horizontal=horizontal),) transformation += (SCCRevectorTransformation(horizontal=horizontal),) transformation += (SCCAnnotateTransformation( @@ -572,7 +572,10 @@ def test_scc_hoist_multiple_kernels_loops(frontend, horizontal, vertical, blocki assert assign.lhs == 'end' assert assign.rhs == 'nlon-nb' assigns = FindNodes(Assignment).visit(driver_loops[1]) - assert not assign in assigns + if trim_vector_sections: + assert not assign in assigns + else: + assert assign in assigns assert driver_loops[4] in FindNodes(Loop).visit(driver_loops[3].body) assert driver_loops[5] in FindNodes(Loop).visit(driver_loops[3].body) diff --git a/transformations/transformations/single_column_coalesced_vector.py b/transformations/transformations/single_column_coalesced_vector.py index db99e4d5a..446e7435d 100644 --- a/transformations/transformations/single_column_coalesced_vector.py +++ b/transformations/transformations/single_column_coalesced_vector.py @@ -228,7 +228,7 @@ def process_driver(self, routine, targets=()): new_driver_loop = loop.clone(body=new_driver_loop) sections = self.extract_vector_sections(new_driver_loop.body, self.horizontal) if self.trim_vector_sections: - sections = self.get_trimmed_sections(routine, self.horizontal, sections) + sections = self.get_trimmed_sections(new_driver_loop, self.horizontal, sections) section_mapper = {s: ir.Section(body=s, label='vector_section') for s in sections} new_driver_loop = NestedTransformer(section_mapper).visit(new_driver_loop) driver_loop_map[loop] = new_driver_loop