Fix a bug in exiting SSA form for loops (#1293)

The Slang compiler was bit by a known issue when translating from SSA form back to straight-line code. Give code like the following:

    int x = 0;
    int y = 1;
    while(...)
    {
        ...
        int t = x;
        x = y;
        y = t;
    }
    ...

The SSA construction pass will eliminate the temporary `t` and yield code something like:

        br(b, 0, 1);

    block b(param x : Int, param y : Int):
        ...
        br(b, y, x);

The loop-dependent variables have become parameters of the loop block, and the branchs to the top of the loop pass the appropriate values for the next iteration (e.g., the jump that starts the loop sends in `0` and `1`).

The problem comes up when translating the back-edge the continues the loop out of SSA form. Our generated code will re-introduce temporaries for `x` and `y`:

    int x;
    int y;

    // jump into loop becomes:
    x = 0;
    y = 1;

    for(;;)
    {
        ...
        // back-edge becomes
        x = y;
        y = x;
        continue;
    }

The problem there is that we've naively translated a branch like `br(b, <a>, <b>)` into `x = <a>; y = <b>;` but that doesn't work correctly in the case where `<b>` is `x`, because we will have already clobbered the value of `x` with `<a>`.

The simplest fix is to introduce a temporary (just like the input code had), and generate:

    // back-edge becomes
    int t = x;
    x = y;
    y = t;

This change modifies the `emitPhiVarAssignments()` function so that it detects bad cases like the above and emits temporaries to work around the problem. A new test case is included that produced incorrect output before the change, and now produces the expected results.

A secondary change is folded in here that tries to guard against a more subtle version of the problem:

    for(...)
    {
        ...
        int x1 = x + 1;
        int y1 = y + 1;
        x = y1;
        y = x1;
    }

In this more complicated case, each of `x` and `y` is being assigned to a value derived from the other, but neither is being set using a block parameter directly, so the changes to `emitPhiVarAssignments()` do not apply.

The problem in this case would be if the `shouldFoldInstIntoUseSites()` logic decided to fold the computation of `x1` or `y1` into the branch instruction, resulting in:

    x = y + 1;
    y = x + 1;

which would again violate the semantics of the original code, because now there is an assignment to `x` before the computation of `x + 1`.

Right now it seems impossible to force this case to arise in practice, due to implementation details in how we generate IR code for loops. In particular, the block that computes the `x+1` and `y+1` values is currently always distinct from the block that branches back to the top of the loop, and we do not allow "folding" of sub-expressions from different blocks. It is possible, however, that future changes to the compiler could change the form of the IR we generate and make it possible for this problem to arise.

The right fix for this issue would be to say that we should introduce a temporary for any branch argument that "involves" a block parameter (whether directly using it or using it as a sub-expression). Unfortunately, the ad hoc approach we use for folding sub-expressions today means that testing if an operand "involves" something would be both expensive and unwieldy.

A more expedient fix is to disallow *all* folding of sub-expressions into unconditional branch instructions (the ones that can pass arguments to the target block), which is what I ended up implementing in this change. Making that defensive change alters the GLSL we output for some of our cross-compilation tests, in a way that required me to update the baseline/gold GLSL.

A better long-term fix for this whole space of issues would be to have the "de-SSA" operation be something we do explicitly on the IR. Such an IR pass would still need to be careful about the first issue addressed in this change, but the second one should (in principle) be a non-issue given that our emit/folding logic already handles code with explicit mutable local variables correctly.

1 files changed, 201 insertions, 9 deletions

diff --git a/source/slang/slang-emit-c-like.cpp b/source/slang/slang-emit-c-like.cpp
index b217a2a1b..bf1754a32 100644
--- a/source/slang/slang-emit-c-like.cpp
+++ b/source/slang/slang-emit-c-like.cpp
@@ -1094,6 +1094,17 @@ bool CLikeSourceEmitter::shouldFoldInstIntoUseSites(IRInst* inst)
             return false;
     }
 
+    // As a safeguard, we should not allow an instruction that references
+    // a block parameter to be folded into a unconcditonal branch
+    // (which includes arguments for the parameters of the target block).
+    //
+    // For simplicity, we will just disallow folding of intructions
+    // into an unconditonal branch completely, and leave a more refined
+    // version of this check for later.
+    //
+    if(as<IRUnconditionalBranch>(user))
+        return false;
+
     // Okay, if we reach this point then the user comes later in
     // the same block, and there are no instructions with side
     // effects in between, so it seems safe to fold things in.
@@ -2406,25 +2417,206 @@ void CLikeSourceEmitter::emitLayoutSemantics(IRInst* inst, char const* uniformSe
 
 void CLikeSourceEmitter::emitPhiVarAssignments(UInt argCount, IRUse* args, IRBlock* targetBlock)
 {
-    UInt argCounter = 0;
-    for (auto pp = targetBlock->getFirstParam(); pp; pp = pp->getNextParam())
+    // The basic setup here is that we are at the end of a block
+    // and are about to branch to `targetBlock`. The target block
+    // has parameters (representing the phi variables/nodes in SSA
+    // form), and the block we are branching from provided arguments
+    // (given as `args` and `argCount` here) to be passed to those
+    // parameters.
+    //
+    // An earlier step will have already emitted a local variable
+    // declaration for each phi node (block parameter), so we should
+    // be able to "pass" an argument to a parameter by assigning
+    // to the variable that represents the parameter.
+    //
+    // A naive approach would simply loop over the parameters/arguments
+    // in tandem and emit assignments like:
+    //
+    //      <param_i> = <arg_i>;
+    //
+    // This approach has an important and known failure case, which
+    // can occur when one of the arguments is also one of the parameters.
+    //
+    // If we have a block like:
+    //
+    //      block b(
+    //          param x: Int,
+    //          y: Int)
+    //          ...
+    //
+    // and then a branch to it like:
+    //
+    //      br(b, y, x);
+    //
+    // Then the naive approach for emitting the assignments for
+    // that branch would output:
+    //
+    //      x = y;
+    //      y = x;
+    //
+    // But that is not the semantics that the "parameter passing"
+    // model is meant to capture.
+    //
+    // The simplest way to restore the correct semantics would
+    // be to assign each of the arguments to a temporary, and
+    // then to assign those temporaries to the parameters:
+    //
+    //      let tmp0 = y;
+    //      let tmp1 = x;
+    //      x = tmp0;
+    //      y = tmp1;
+    //
+    // Adding temporaries like that unconditionally would clutter
+    // up the generated code, so it would be nicer to only
+    // introduce them when strictly necessary.
+    //
+    // A temporary should only be necessary when we have:
+    //
+    // * A parameter `destParam` of the `targetBlock`
+    // * Where the argument value for `destParam` is another parameter, `srcParam`, of `targetBlock`
+    // * And `srcParam` appears before `destParam` in the parameter list
+    //
+    // We start by looking for such cases, and collecting the
+    // set of block parameters that will require a defensive
+    // temporary.
+    //
+    // Note: an alternative approach would first look for a total
+    // order on the arguments/parameters such that we avoid
+    // the problem, and then only resort to introducing temporaries
+    // as a means of breaking cycles.
+    //
+    // We will track the parameters that need temporaries by
+    // building a dictionary mapping parameters to the name
+    // of the temporary they will use, if any.
+    //
+    Dictionary<IRParam*, String> mapParamToTempName;
     {
-        UInt argIndex = argCounter++;
+        // To know whether a  parameter occurs earlier in
+        // the list than another, we will rely on a set
+        // of parameters we have already "seen" during
+        // our iteration.
+        //
+        // TODO: This is a strong candidate for using a
+        // hash set optimized for a small number of entries
+        // (and thus using a stack allocation in the common case),
+        // given that most blocks will only have a handful
+        // of parameters.
+        //
+        HashSet<IRParam*> seenParams;
 
-        if (argIndex >= argCount)
+        UInt argCounter = 0;
+        for (auto destParam = targetBlock->getFirstParam(); destParam; destParam = destParam->getNextParam())
         {
-            SLANG_UNEXPECTED("not enough arguments for branch");
-            break;
+            seenParams.Add(destParam);
+
+            UInt argIndex = argCounter++;
+            if (argIndex >= argCount)
+            {
+                SLANG_UNEXPECTED("not enough arguments for branch");
+                break;
+            }
+            IRInst* arg = args[argIndex].get();
+
+            // Is the argument also a parameter of the same block?
+            //
+            // If not, then we don't have to worry about this assignment.
+            //
+            IRParam* srcParam = as<IRParam>(arg);
+            if(!srcParam)
+                continue;
+            if(srcParam->getParent() != targetBlock)
+                continue;
+
+            // Is the argument an *earlier* parameter of the same block?
+            // We test this by checking if we've already seen it during
+            // our iteration.
+            //
+            // If the parameter isn't an earlier one, then we don't have
+            // to worry about the order of assignment creating problems.
+            //
+            if(!seenParams.Contains(srcParam))
+                continue;
+
+            // At this point we've detected a problem case: the `dstParam`
+            // would naively be assigned the value of `srcParam` *after*
+            // we've already assigned to `srcParam`.
+            //
+            // We will need to avoid the problem by introducing a temporary
+            // for `srcParam`.
+            //
+            if( !mapParamToTempName.ContainsKey(srcParam) )
+            {
+                String tempName = "tmp_" + getName(srcParam);
+                mapParamToTempName.Add(srcParam, tempName);
+            }
         }
+    }
 
-        IRInst* arg = args[argIndex].get();
+    // Now that we've determined which of the parameters could cause problems,
+    // we can do an initial round of assignments, where we move form the
+    // argument values into either a parameter (if there were no problems)
+    // or a temporary (if we needed to avoid a problem).
+    //
+    {
+        UInt argCounter = 0;
+        for (auto param = targetBlock->getFirstParam(); param; param = param->getNextParam())
+        {
+            UInt argIndex = argCounter++;
+            if (argIndex >= argCount)
+            {
+                SLANG_UNEXPECTED("not enough arguments for branch");
+                break;
+            }
+            IRInst* arg = args[argIndex].get();
+
+            auto outerPrec = getInfo(EmitOp::General);
+            auto prec = getInfo(EmitOp::Assign);
+
+            String tempName;
+            if( mapParamToTempName.TryGetValue(param, tempName) )
+            {
+                // This parameter was involved in a confounding assignment,
+                // where it was used as the argument value for a later
+                // parameter on this same branch/edge.
+                //
+                // We need to assign to a temporary instead of
+                // to the parameter itself. We will declare
+                // the temporary here so that the assignment
+                // of the argument serves as the initializer.
+                //
+                emitType(param->getFullType(), tempName);
+            }
+            else
+            {
+                // This parameter was not involved in any confounding assignments,
+                // so we can simply us the parameter itself as the left-hand side
+                // of an assignment.
+                //
+                emitOperand(param, leftSide(outerPrec, prec));
+            }
+
+            m_writer->emit(" = ");
+            emitOperand(arg, rightSide(prec, outerPrec));
+            m_writer->emit(";\n");
+        }
+    }
+
+    // Finally, after all the assignments from argument values have been completed,
+    // we will make a cleanup pass that copies from any of the temporary variables
+    // we introduced over to the parameter that needed a temporary.
+    //
+    for (auto param = targetBlock->getFirstParam(); param; param = param->getNextParam())
+    {
+        String tempName;
+        if( !mapParamToTempName.TryGetValue(param, tempName) )
+            continue;
 
         auto outerPrec = getInfo(EmitOp::General);
         auto prec = getInfo(EmitOp::Assign);
 
-        emitOperand(pp, leftSide(outerPrec, prec));
+        emitOperand(param, leftSide(outerPrec, prec));
         m_writer->emit(" = ");
-        emitOperand(arg, rightSide(prec, outerPrec));
+        m_writer->emit(tempName);
         m_writer->emit(";\n");
     }
 }