stfc / PSyclone

Domain-specific compiler and code transformation system for Finite Difference/Volume/Element Earth-system models in Fortran
BSD 3-Clause "New" or "Revised" License
107 stars 29 forks source link

New transformation to replace a single-trip loop with an assignment to the loop variable #2738

Open arporter opened 1 month ago

arporter commented 1 month ago

Given the reshaping of arrays being passed into the UM physics routines, there are many cases where there are loops over a dimension that we know, at build time, will have an extent of just unity. To help the compiler, it would therefore be good to remove all such loops and simply replace them with an assignment to the loop variable followed by the body of the loop:

        if loop.variable.name == "jj":
            # Ensure any symbols that are in the table of the Loop are moved
            # into the outer scope.
            parent_table = loop.parent.scope.symbol_table
            parent_table.merge(loop.loop_body.symbol_table)

            # Move the body of the loop after the loop
            for statement in reversed(loop.loop_body.children):
                loop.parent.addchild(statement.detach(), loop.position + 1)

            # Create the Assignment node
            assign = Assignment.create(Reference(loop.variable),
                                       Literal("1", INTEGER_TYPE))

            # Replace the (now empty) Loop with the Assignment
            loop.replace_with(assign)

Since this isn't entirely trivial, it would be good to have this as a transformation. I'm not sure what it should be called though. Perhaps ReplaceSingleTripLoopTrans?

arporter commented 1 month ago

Since most of the implementation is given in the issue description, this could be a good, self-contained issue for a new person to tackle.

arporter commented 1 month ago

Note, if the body of the loop doesn't contain any CodeBlocks then we could go one better and replace all References to the loop variable with the value supplied to the transformation, i.e. "1" in the example given above.

hiker commented 1 month ago

I believe we recently discussed to generalising the replace_induction_variables_trans? That would be another use case for a more generalised approach (the replace_induction... transformation just needs to restrict the constant replacement to the loop body)

arporter commented 1 month ago

I believe we recently discussed to generalising the replace_induction_variables_trans? That would be another use case for a more generalised approach (the replace_induction... transformation just needs to restrict the constant replacement to the loop body)

Oh yes, can you remember whether we had that discussion on GH and, if so, where? Do you mean that we'd use ReplaceInductionVariablesTrans to handle the substitution of the loop variable within the body?

hiker commented 1 month ago

It was on teams, see message there. My feeling is that we have a more generic constant replacement transformation (without thinking too much: a list of sibling nodes, where a variable is written once and then only read. We can then replace the value (and need to figure out what to do with the rest of the code that's not part of the replacement region). And loop induction variable would just do some additional validation(?) and call the more generic function with just the loop body?

hiker commented 1 month ago

Here my comment on teams (which might be a bit long):

During development of the trainings material, I came across the following loop structures after inlining (code using dl_ems_inf, but not gocean, so it's transformation only):

xstart = current%internal%xstart
xstop = current%internal%xstop
ystart = current%internal%ystart
ystop = current%internal%ystop
xstart_1 = current%internal%xstart
xstop_1 = current%internal%xstop
ystart_1 = current%internal%ystart
ystop_1 = current%internal%ystop
do j = ystart, ystop, 1
do i = xstart, xstop, 1
neighbours%data(i,j) = current%data(i - 1,j - 1) + current%data(i,j - 1) + ...
enddo
enddo
do j_1 = ystart_1, ystop_1, 1
do i_1 = xstart_1, xstop_1, 1
born%data(i_1,j_1) = 0.0
if (current%data(i_1,j_1) == 0.0 .AND. neighbours%data(i_1,j_1) == 3.0) then
born%data(i_1,j_1) = 1.0
end if
enddo
enddo

These loops cannot be fused, since the boundaries appear to be different (even setting the force or what's it called option didn't work, force only disables the dep-analysis., not the loop boundary test). While I have fixed this by explicitly replacing the variables with the values (and then I could fuse the loops), I am wondering: our ReplaceInductionVariable transformation nearly does this already, except it tests that it is inside a loop etc. Would it make sense to make this transformation more generic as a 'replace constant values in a basis block' transformation? Then either make the induction-replacement derive from this with additional validation, or even replace it entirely?