From 5dbb3b4b01e33e96a78d0214cc492a4930e23d67 Mon Sep 17 00:00:00 2001 From: Rolf Heilemann Myhre Date: Fri, 24 Nov 2023 12:17:55 +0100 Subject: [PATCH 1/4] pool_allocator handles range indices --- transformations/tests/test_pool_allocator.py | 8 ++++---- transformations/transformations/pool_allocator.py | 5 ++++- 2 files changed, 8 insertions(+), 5 deletions(-) diff --git a/transformations/tests/test_pool_allocator.py b/transformations/tests/test_pool_allocator.py index b5c0340b5..29327abab 100644 --- a/transformations/tests/test_pool_allocator.py +++ b/transformations/tests/test_pool_allocator.py @@ -454,7 +454,7 @@ def test_pool_allocator_temporaries_kernel_sequence(frontend, block_dim, directi integer, parameter :: jprb = selected_real_kind(13,300) integer, intent(in) :: start, end, klon, klev real(kind=jprb), intent(inout) :: field2(klon,klev) - real(kind=jprb) :: tmp1(2*klon, klev), tmp2(klon, klev) + real(kind=jprb) :: tmp1(2*klon, klev), tmp2(klon, 0:klev) integer :: jk, jl do jk=1,klev @@ -520,7 +520,7 @@ def test_pool_allocator_temporaries_kernel_sequence(frontend, block_dim, directi # pylint: disable-next=line-too-long assert kernel_item.trafo_data[transformation._key]['stack_size'] == f'c_sizeof(real(1, kind={kind_real}))*klon + c_sizeof(real(1, kind={kind_real}))*klev*klon + 2*c_sizeof(int(1, kind={kind_int}))*klon + c_sizeof(logical(true, kind={kind_log}))*klev' # pylint: disable-next=line-too-long - assert kernel2_item.trafo_data[transformation._key]['stack_size'] == f'3*c_sizeof(real(1, kind={kind_real}))*klev*klon' + assert kernel2_item.trafo_data[transformation._key]['stack_size'] == f'3*c_sizeof(real(1, kind={kind_real}))*klev*klon + c_sizeof(real(1, kind={kind_real}))*klon' assert all(v.scope is None for v in FindVariables().visit(kernel_item.trafo_data[transformation._key]['stack_size'])) assert all(v.scope is None for v in @@ -549,7 +549,7 @@ def test_pool_allocator_temporaries_kernel_sequence(frontend, block_dim, directi stack_size = f'max(c_sizeof(real(1, kind={kind_real}))*nlon + c_sizeof(real(1, kind={kind_real}))*nlon*nz + ' stack_size += f'2*c_sizeof(int(1, kind={kind_int}))*nlon + c_sizeof(logical(true, kind={kind_log}))*nz,' - stack_size += f'3*c_sizeof(real(1, kind={kind_real}))*nz*nlon)/c_sizeof(real(1, kind=jprb))' + stack_size += f'3*c_sizeof(real(1, kind={kind_real}))*nlon*nz+c_sizeof(real(1,kind=jprb))*nlon)/c_sizeof(real(1, kind=jprb))' check_stack_created_in_driver(driver, stack_size, calls[0], 2) # Has the data sharing been updated? @@ -617,7 +617,7 @@ def test_pool_allocator_temporaries_kernel_sequence(frontend, block_dim, directi elif ass.lhs == 'ylstack%l' and 'ylstack%l' in ass.rhs and 'c_sizeof' in ass.rhs and dim1 == _size: # Stack increment for tmp1 assign_idx['tmp1_stack_incr'] = idx - elif ass.lhs == 'ylstack%l' and 'ylstack%l' in ass.rhs and 'c_sizeof' in ass.rhs and dim2 == _size: + elif ass.lhs == 'ylstack%l' and 'ylstack%l' in ass.rhs and 'c_sizeof' in ass.rhs: # Stack increment for tmp2 assign_idx['tmp2_stack_incr'] = idx diff --git a/transformations/transformations/pool_allocator.py b/transformations/transformations/pool_allocator.py index d43eaf2de..a044efaff 100644 --- a/transformations/transformations/pool_allocator.py +++ b/transformations/transformations/pool_allocator.py @@ -477,7 +477,10 @@ def _create_stack_allocation(self, stack_ptr, stack_end, ptr_var, arr, stack_siz # Build expression for array size in bytes dim = arr.dimensions[0] for d in arr.dimensions[1:]: - dim = Product((dim, d)) + if isinstance(d, RangeIndex): + dim = simplify(Product((dim, Sum((d.upper, Product((-1,d.lower)), IntLiteral(1)))))) + else: + dim = Product((dim, d)) arr_size = Product((dim, InlineCall(Variable(name='C_SIZEOF'), parameters=as_tuple(self._get_c_sizeof_arg(arr))))) From c1a75a2b700c7e46622d4da6a5d22022b1f54eb0 Mon Sep 17 00:00:00 2001 From: Rolf Heilemann Myhre Date: Fri, 24 Nov 2023 12:32:38 +0100 Subject: [PATCH 2/4] split long line --- transformations/tests/test_pool_allocator.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/transformations/tests/test_pool_allocator.py b/transformations/tests/test_pool_allocator.py index 29327abab..297bff942 100644 --- a/transformations/tests/test_pool_allocator.py +++ b/transformations/tests/test_pool_allocator.py @@ -549,7 +549,8 @@ def test_pool_allocator_temporaries_kernel_sequence(frontend, block_dim, directi stack_size = f'max(c_sizeof(real(1, kind={kind_real}))*nlon + c_sizeof(real(1, kind={kind_real}))*nlon*nz + ' stack_size += f'2*c_sizeof(int(1, kind={kind_int}))*nlon + c_sizeof(logical(true, kind={kind_log}))*nz,' - stack_size += f'3*c_sizeof(real(1, kind={kind_real}))*nlon*nz+c_sizeof(real(1,kind=jprb))*nlon)/c_sizeof(real(1, kind=jprb))' + stack_size += f'3*c_sizeof(real(1, kind={kind_real}))*nlon*nz + ' + stack_size += f'c_sizeof(real(1,kind=jprb))*nlon)/c_sizeof(real(1, kind=jprb))' check_stack_created_in_driver(driver, stack_size, calls[0], 2) # Has the data sharing been updated? From 036153a189abe6ace6814238adaa4249ed1eafc4 Mon Sep 17 00:00:00 2001 From: Rolf Heilemann Myhre Date: Fri, 24 Nov 2023 12:39:14 +0100 Subject: [PATCH 3/4] f string stuff --- transformations/tests/test_pool_allocator.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/transformations/tests/test_pool_allocator.py b/transformations/tests/test_pool_allocator.py index 297bff942..85c7769ce 100644 --- a/transformations/tests/test_pool_allocator.py +++ b/transformations/tests/test_pool_allocator.py @@ -550,7 +550,7 @@ def test_pool_allocator_temporaries_kernel_sequence(frontend, block_dim, directi stack_size = f'max(c_sizeof(real(1, kind={kind_real}))*nlon + c_sizeof(real(1, kind={kind_real}))*nlon*nz + ' stack_size += f'2*c_sizeof(int(1, kind={kind_int}))*nlon + c_sizeof(logical(true, kind={kind_log}))*nz,' stack_size += f'3*c_sizeof(real(1, kind={kind_real}))*nlon*nz + ' - stack_size += f'c_sizeof(real(1,kind=jprb))*nlon)/c_sizeof(real(1, kind=jprb))' + stack_size += f'c_sizeof(real(1,kind={kind_real}))*nlon)/c_sizeof(real(1, kind=jprb))' check_stack_created_in_driver(driver, stack_size, calls[0], 2) # Has the data sharing been updated? From 4a37b2dd645b0f479e8c311d7ddde977b3085cd2 Mon Sep 17 00:00:00 2001 From: Rolf Heilemann Myhre Date: Fri, 24 Nov 2023 13:05:33 +0100 Subject: [PATCH 4/4] missing space making test freak out? --- transformations/tests/test_pool_allocator.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/transformations/tests/test_pool_allocator.py b/transformations/tests/test_pool_allocator.py index 85c7769ce..73bd365df 100644 --- a/transformations/tests/test_pool_allocator.py +++ b/transformations/tests/test_pool_allocator.py @@ -550,7 +550,7 @@ def test_pool_allocator_temporaries_kernel_sequence(frontend, block_dim, directi stack_size = f'max(c_sizeof(real(1, kind={kind_real}))*nlon + c_sizeof(real(1, kind={kind_real}))*nlon*nz + ' stack_size += f'2*c_sizeof(int(1, kind={kind_int}))*nlon + c_sizeof(logical(true, kind={kind_log}))*nz,' stack_size += f'3*c_sizeof(real(1, kind={kind_real}))*nlon*nz + ' - stack_size += f'c_sizeof(real(1,kind={kind_real}))*nlon)/c_sizeof(real(1, kind=jprb))' + stack_size += f'c_sizeof(real(1, kind={kind_real}))*nlon)/c_sizeof(real(1, kind=jprb))' check_stack_created_in_driver(driver, stack_size, calls[0], 2) # Has the data sharing been updated?