Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

pool_allocator handles range indices #193

Merged
merged 5 commits into from
Nov 29, 2023
Merged

pool_allocator handles range indices #193

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
5dbb3b4
pool_allocator handles range indices
rolfhm Nov 24, 2023
c1a75a2
split long line
rolfhm Nov 24, 2023
036153a
f string stuff
rolfhm Nov 24, 2023
4a37b2d
missing space making test freak out?
rolfhm Nov 24, 2023
bc5e4e7
Merge branch 'main' into sbrm-range-indices-in-stack
rolfhm Nov 27, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 5 additions & 4 deletions transformations/tests/test_pool_allocator.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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}))*nz*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={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?
Expand Down Expand Up @@ -617,7 +618,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

Expand Down
5 changes: 4 additions & 1 deletion transformations/transformations/pool_allocator.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -481,7 +481,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)))))

Expand Down
Loading