Comment by James Molloy (JIRA)

This still reproduces at HEAD today. Mingsheng, are you looking at this or are you OK for me to steal it? 🙂

Comment by Mingsheng Hong (JIRA)

James, please feel free to take it. Thanks for the help!

Looks like at the end of TFGraphLowering::visitBuiltinRecvFromHostInst(), we need to call addValueMapping() to map the output of the "__tfop_tfc.RecvFromHost" builtin to a graph node that it constructed, just like in addTFRecvOp().

Happy to support as needed.

Comment by Mingsheng Hong (JIRA)

Correction: We do call addValueMapping() at

When debugging another issue with Chris, we found that the hack of early return here based on shouldLowerEffectfulOps is incorrect

https://github.com/apple/swift/blob/553ac93bb98f08adc2f88b2ce76b245a92a7b5ab/lib/SILOptimizer/Mandatory/TFLowerGraph.cpp#L963. When TF needs to recv a value from Swift host within a loop, that recv is currently present in both the loop header and body. In that case we cannot simply do early return when lowering the recv node in the loop header (where

shouldLowerEffectfulOps is set for loop header at https://github.com/apple/swift/blob/553ac93bb98f08adc2f88b2ce76b245a92a7b5ab/lib/SILOptimizer/Mandatory/TFLowerGraph.cpp#L2076),). Instead we'll need to issue an internal error, and wait for proper loop rotation to fix this kind of issues.

That should be the same root cause for the above test case, where the output scalar of
(a % 2).elementsEqual(0).scalar!
, computed on the host, is sent back to TF.

Comment by Mingsheng Hong (JIRA)

#17502 fixed a bug in the current hacky solution – will need to wait for loop rotation to land as a proper fix!

CC jmolloy (JIRA User) bgogul (JIRA User)

Comment by James Molloy (JIRA)

To be clear, I don't think what we need here is classic loop rotation, I think it's a little simpler than that.

Canonical loop form is head-tested:

Preheader
    |
Header <+
    |   |
Latch --+
    |
 Exit

Loop rotated form is tail-tested:

Preheader
 |    |
 |  Header <+
 |    |     |
 |  Latch --+
 |    |
 +> Exit

We want a head-tested loop (While form, not Do-While); we just want the header to contain only the loop test and no other code. We can do this by pushing (and duplicating) all instructions except the condbr from the header into both preheader and latch.

This is a very simple transform that should only be a few lines of code as part of TFCanonicalizeCFG, rather than full-on loop rotation which is a bigger problem.

James

Comment by Gogul Balakrishnan (JIRA)

James, you are right. It does not require full loop rotation. My fix for SR-7765 does exactly as you had suggested—canonicalizes the loop so that only side-effect free operations are in the header.

Comment by Mingsheng Hong (JIRA)

pushing (and duplicating) all instructions except the condbr from the header into both preheader and latch

When the loop condition has a side-effecting op, it seems the above is not sufficient?

e.g. if the code has:

 var a = <some tensor>
 while (A < SomeTensorFromHost()) {
   a += ...
 }
 

Compiler will need to transform the code into:

 var a = <some tensor>
 var cond = A < SomeTensorFromHost()
 while (cond) {
   a += ...
   cond = A < SomeTensorFromHost()
 }

Is this transformation usually referred to as loop rotation?

Comment by James Molloy (JIRA)

Kind of; loop rotation normally involves changing a loop from being head-tested to tail-tested so the loop test can use state calculated in the loop body and only one phi is needed, for the induction variable.

In this case we still want the loop to be head-tested, so it's not loop rotation per se. But that's just semantics (I posted it to explain why we couldn't just use the SIL loop rotation pass that already exists).