Hoist transformation with multiple kernels allowing hoisted variable of a duplicated name but different size.

[Zbigniew Piotrowski]

When two or more kernels are called from a driver, hoisting machinery works largely independently on each kernel. If the hoisted variable in kernel 1 has the same name but different size than in kernel 2, this results in generation of invalid Fortran code (duplicated hoisted variable declaration). The following example detects such situation and raises an error in case there are two variables of the same name but of different size and/or dimensions. It is desired that either:
the transformation stops, or
the conflicting hoisted variables are renamed to e.g. indicate their kernel of origin.

# This software is licensed under the terms of the Apache Licence Version 2.0
# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
# In applying this licence, ECMWF does not waive the privileges and immunities
# granted to it by virtue of its status as an intergovernmental organisation
# nor does it submit to any jurisdiction.

import pytest

from loki import (
    Subroutine, Dimension, fgen, Sourcefile, SubroutineItem
)
from conftest import available_frontends
from transformations import (
    SCCDevectorTransformation, SCCHoistTransformation
)
#pylint: disable=too-many-lines

@pytest.fixture(scope='module', name='horizontal')
def fixture_horizontal():
    return Dimension(name='horizontal', size='nlon', index='jl', bounds=('start', 'end'))


@pytest.fixture(scope='module', name='vertical')
def fixture_vertical():
    return Dimension(name='vertical', size='nz', index='jk')


@pytest.fixture(scope='module', name='blocking')
def fixture_blocking():
    return Dimension(name='blocking', size='nb', index='b')



@pytest.mark.parametrize('frontend', available_frontends())
def test_scc_hoist_multiple_kernels(frontend, horizontal, vertical, blocking):
    """
    Test hoisting of column temporaries to "driver" level.
    """

    fcode_driver = """
  SUBROUTINE column_driver(nlon, nz, q, nb)
    INTEGER, INTENT(IN)   :: nlon, nz, nb  ! Size of the horizontal and vertical
    REAL, INTENT(INOUT)   :: q(nlon,nz,nb)
    INTEGER :: b, start, end

    start = 1
    end = nlon
    do b=1, nb
      call compute_column(start, end, nlon, nz, q(:,:,b))

      call compute_column2(start, end, nlon, nz, q(:,:,b))

      call compute_column3(start, end, nlon, nz, q(:,:,b))
    end do
  END SUBROUTINE column_driver
"""

    fcode_kernel3 = """
  SUBROUTINE compute_column3(start, end, nlon, nz, q)
    INTEGER, INTENT(IN) :: start, end  ! Iteration indices
    INTEGER, INTENT(IN) :: nlon, nz    ! Size of the horizontal and vertical
    REAL, INTENT(INOUT) :: q(nlon,nz)
    REAL :: T(NLON,NZ)
    REAL :: TALT(nlon,nz)
    REAL :: TYETALT(nlon,NZ,1,2,3)
    INTEGER :: jl, jk
    REAL :: c

    c = 5.345
    DO jk = 2, nz
      DO jl = start, end
        t(jl, jk) = c * k
        q(jl, jk) = q(jl, jk-1) + t(jl, jk) * c
      END DO
    END DO

    ! The scaling is purposefully upper-cased
    DO JL = START, END
      Q(JL, NZ) = Q(JL, NZ) * C
    END DO
  END SUBROUTINE compute_column3
"""
    fcode_kernel2 = """
  SUBROUTINE compute_column2(start, end, nlon, nz, q)
    INTEGER, INTENT(IN) :: start, end  ! Iteration indices
    INTEGER, INTENT(IN) :: nlon, nz    ! Size of the horizontal and vertical
    REAL, INTENT(INOUT) :: q(nlon,nz)
    REAL :: t(nlon,nz,2)
    REAL :: TALT(nlon,nz)
    REAL :: tyetalt(NLON,nz,1,2,3)
    INTEGER :: jl, jk
    REAL :: c

    c = 5.345
    DO jk = 2, nz
      DO jl = start, end
        t(jl, jk, 2) = c * k
        q(jl, jk) = q(jl, jk-1) + t(jl, jk, 2) * c
      END DO
    END DO

    ! The scaling is purposefully upper-cased
    DO JL = START, END
      Q(JL, NZ) = Q(JL, NZ) * C
    END DO
  END SUBROUTINE compute_column2
"""
    fcode_kernel = """
  SUBROUTINE compute_column(start, end, nlon, nz, q)
    INTEGER, INTENT(IN) :: start, end  ! Iteration indices
    INTEGER, INTENT(IN) :: nlon, nz    ! Size of the horizontal and vertical
    REAL, INTENT(INOUT) :: q(nlon,nz)
    REAL :: t(nlon,nz)
    REAL :: talt(nlon,nz)
    REAL :: tyetalt(nlon,nz,1)
    INTEGER :: jl, jk
    REAL :: c

    c = 5.345
    DO jk = 2, nz
      DO jl = start, end
        t(jl, jk) = c * k
        q(jl, jk) = q(jl, jk-1) + t(jl, jk) * c
      END DO
    END DO

    ! The scaling is purposefully upper-cased
    DO JL = START, END
      Q(JL, NZ) = Q(JL, NZ) * C
    END DO
  END SUBROUTINE compute_column
"""
    driver_item = SubroutineItem(name='#column_driver',
                                 source=Sourcefile.from_source(fcode_driver, frontend=frontend))
    kernel = Subroutine.from_source(fcode_kernel, frontend=frontend)
    kernel2 = Subroutine.from_source(fcode_kernel2, frontend=frontend)
    kernel3 = Subroutine.from_source(fcode_kernel3, frontend=frontend)
    driver = Subroutine.from_source(fcode_driver, frontend=frontend)
    driver.enrich_calls(kernel)  # Attach kernel source to driver call
    driver.enrich_calls(kernel2)  # Attach kernel source to driver call
    driver.enrich_calls(kernel3)  # Attach kernel source to driver call

    scc_transform = (SCCDevectorTransformation(horizontal=horizontal),)
    scc_transform += (SCCHoistTransformation(horizontal=horizontal, vertical=vertical,
                                             block_dim=blocking),)
    for transform in scc_transform:
        transform.apply(driver, role='driver',
                        targets=['compute_column', 'compute_column2', 'compute_column3'],
                        item=driver_item)
        transform.apply(kernel, role='kernel')
        transform.apply(kernel2, role='kernel')
        transform.apply(kernel3, role='kernel')
    hoisted_vars = driver_item.trafo_data['SCCHoistTransformation']['column_locals']
    varnames = []
    varnamesdims = []
#Extract hoisted variable names and dimensions
    for i, hoist_var in enumerate(hoisted_vars):
        varnames.append(hoist_var.name.lower())
        varnamesdims.append(hoist_var.dimensions)

#Check if there are duplicated v
    if len(varnames) != len(set(varnames)):
        vardup = []
        vardupt = []
        varsng = []
        varsngt = []
        for i, hoisted_name in enumerate(varnames):
            if hoisted_name not in varsng:
                varsng.append(hoisted_name)
                varsngt.append((hoisted_name, i, varnamesdims[i]))
            else:
                vardup.append(hoisted_name)
                vardupt.append((hoisted_name, i, varnamesdims[i]))

#check if duplicated variables have dimensions of different size
    for i, duplicated_name in enumerate(vardup):
        j = varsng.index(duplicated_name)
        vdims1 = fgen(varsngt[j][2]).lower()
        vdims2 = fgen(vardupt[j][2]).lower()
        if vdims1 != vdims2:
            print("First   occurence of hoisted variable: ", duplicated_name,
                  " has dimensions: ", ' '.join(vdims1.split()))
            print("Another occurence of hoisted variable: ", duplicated_name,
                  " has dimensions: ", ' '.join(vdims2.split()))

#Return on first duplication
        assert vdims1 == vdims2

[Michael Lange]

This is trivially solved by PR #156, as the recursive hoisting utility always prepends the defining child routines' name to the hoisted variables, thus naturally avoiding variable shadowing. I've locally adopted this example test to confirm that, but it seems redundant, as we should be testing this without the SCC wrapper in a separate test already. I leave it to the integrator of #156 to shout if this should be added for coverage (@reuterbal).