OpenMP collapse=[int] not working for directly nested loops with directives

[schreiberx]

Using the OpenMP directive transformations automatically changes the provided integer value of collapse if do loops are not directly nested but are prefixed or encapsulated with a Directive.

However, if there are other Directives in one of the nested loops, this raises GenerationError:

        # If there is a collapse clause, there must be as many immediately
        # nested loops as the collapse value
        if self._collapse:
            cursor = self.dir_body.children[0]
            for depth in range(self._collapse):
                if not isinstance(cursor, Loop):
                    raise GenerationError(
                        f"OMPLoopDirective must have as many immediately "
                        f"nested loops as the collapse clause specifies but "
                        f"'{self}' has a collapse={self._collapse} and the "
                        f"nested statement at depth {depth} is a "
                        f"{type(cursor).__name__} rather than a Loop.")
                cursor = cursor.loop_body.children[0]

[LonelyCat124]

I think we'd discussed this before internally, and decided we can't support anything other than directly nested loops until we're sure about compiler support as its a newer (5.0?) OpenMP feature. Can you put a replicator of the issue for us to test this in the future?

[schreiberx]

Here we go:

from psyclone.psyir.frontend.fortran import FortranReader
from psyclone.psyir.backend.fortran import FortranWriter
from psyclone.psyir.nodes.omp_directives import OMPSimdDirective
from psyclone.psyir.transformations.omp_loop_trans import OMPLoopTrans


def test_collapse_problem_reproducer():
    """Test that the problem described in issue #1937 is resolved."""
    source = """
subroutine foo()
    integer :: i, j
    real :: a

    do i=1,10
      do j=1,10
        a = 1.0
      end do
    end do

end subroutine foo
"""

    psyir = FortranReader(
        ignore_comments=False, ignore_directives=False
    ).psyir_from_source(source)

    # Wrap the inner loop with !$omp simd
    inner_loop = psyir.walk(Loop)[1]
    inner_loop_parent = inner_loop.parent
    inner_loop_position = inner_loop.position
    simd = OMPSimdDirective()
    simd.dir_body.addchild(inner_loop.detach())
    inner_loop_parent.addchild(simd, index=inner_loop_position)

    outer_loop = psyir.walk(Loop)[0]
    OMPLoopTrans(omp_directive="paralleldo").apply(outer_loop, collapse=2)

    result = FortranWriter()(psyir)
    print(result)
    assert "collapse(2)" in result


test_collapse_problem_reproducer()

The last assert fails.

[sergisiso]

But we are not certain if this:

!$omp parallel do collapse(2) private(i,j)
do i = 1, 10, 1
  !$omp simd
  do j = 1, 10, 1

is valid (or supported by compilers implementing older versions), for instance, XL Fortran clearly says

The loops must be perfectly nested; that is, there is no intervening code nor any OpenMP directive between the loops which are collapsed.
https://www.ibm.com/docs/en/xl-fortran-linux/16.1.0?topic=clauses-collapse

Although I admit the standard is a bit more ambiguous (but even then I am not sure if it accepts overlapping directives with different semantics)

[sergisiso]

In fact the collapse does a combined iteration space with both loops, I am not sure that then you can attach a property to just part of the iteration space. Would allowing for

!$omp parallel do simd collapse(2) private(i,j)
do i = 1, 10, 1
  do j = 1, 10, 1

work for you? I think this is more clearly std-compliant.

But tbh the current:

!$omp parallel do collapse(1) private(i,j)
do i = 1, 10, 1
  !$omp simd
  do j = 1, 10, 1

it is probably what you want if you are sure you want to vecotrise just the inner loop.

[schreiberx]

Yes, this would also work for me to put the "simd" before the outer loop. It's probably how we'd all do it :-).

The reason to open this issue was solely that the "simd" stopped the collapse(2) from working.
Maybe the automatic parallelization transformation could "move" the simd from the inner loop to the outer one to accommodate this issue (but it's not required for what I'm working on - I can work around this).