diff --git a/transformations/tests/test_single_column_coalesced.py b/transformations/tests/test_single_column_coalesced.py index 619a38f52..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( @@ -567,6 +567,16 @@ 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]) + 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) assert driver_loops[6] in FindNodes(Loop).visit(driver_loops[3].body) @@ -596,6 +606,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) 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