Change the FLATTEN heuristic to "ifs with a loop with a gradient"

src/compiler/translator/ASTMetadataHLSL.cpp

@@ -139,13 +139,16 @@ class PullGradient : public TIntermTraverser
     std::vector<TIntermNode*> mParents;
 };
 
-// Traverses the AST of a function definition, assuming it has already been used to
-// traverse the callees of that function; computes the discontinuous loops and the if
-// statements that contain a discontinuous loop in their call graph.
-class PullComputeDiscontinuousLoops : public TIntermTraverser
+// Traverses the AST of a function definition to compute the the discontinuous loops
+// and the if statements containing gradient loops. It assumes that the gradient loops
+// (loops that contain a gradient) have already been computed and that it has already
+// traversed the current function's callees.
+class PullComputeDiscontinuousAndGradientLoops : public TIntermTraverser
 {
   public:
-    PullComputeDiscontinuousLoops(MetadataList *metadataList, size_t index, const CallDAG &dag)
+    PullComputeDiscontinuousAndGradientLoops(MetadataList *metadataList,
+                                             size_t index,
+                                             const CallDAG &dag)
         : TIntermTraverser(true, false, true),
           mMetadataList(metadataList),
           mMetadata(&(*metadataList)[index]),
@@ -161,15 +164,15 @@ class PullComputeDiscontinuousLoops : public TIntermTraverser
         ASSERT(mIfs.empty());
     }
 
-    // Called when a discontinuous loop or a call to a function with a discontinuous loop
-    // in its call graph is found.
-    void onDiscontinuousLoop()
+    // Called when traversing a gradient loop or a call to a function with a
+    // gradient loop in its call graph.
+    void onGradientLoop()
     {
-        mMetadata->mHasDiscontinuousLoopInCallGraph = true;
+        mMetadata->mHasGradientLoopInCallGraph = true;
         // Mark the latest if as using a discontinuous loop.
         if (!mIfs.empty())
         {
-            mMetadata->mIfsContainingDiscontinuousLoop.insert(mIfs.back());
+            mMetadata->mIfsContainingGradientLoop.insert(mIfs.back());
         }
     }
 
@@ -178,6 +181,11 @@ class PullComputeDiscontinuousLoops : public TIntermTraverser
         if (visit == PreVisit)
         {
             mLoopsAndSwitches.push_back(loop);
+
+            if (mMetadata->hasGradientInCallGraph(loop))
+            {
+                onGradientLoop();
+            }
         }
         else if (visit == PostVisit)
         {
@@ -199,9 +207,9 @@ class PullComputeDiscontinuousLoops : public TIntermTraverser
             ASSERT(mIfs.back() == node);
             mIfs.pop_back();
             // An if using a discontinuous loop means its parents ifs are also discontinuous.
-            if (mMetadata->mIfsContainingDiscontinuousLoop.count(node) > 0 && !mIfs.empty())
+            if (mMetadata->mIfsContainingGradientLoop.count(node) > 0 && !mIfs.empty())
             {
-                mMetadata->mIfsContainingDiscontinuousLoop.insert(mIfs.back());
+                mMetadata->mIfsContainingGradientLoop.insert(mIfs.back());
             }
         }
 
@@ -221,7 +229,6 @@ class PullComputeDiscontinuousLoops : public TIntermTraverser
                     if (loop != nullptr)
                     {
                         mMetadata->mDiscontinuousLoops.insert(loop);
-                        onDiscontinuousLoop();
                     }
                 }
                 break;
@@ -237,7 +244,6 @@ class PullComputeDiscontinuousLoops : public TIntermTraverser
                     }
                     ASSERT(loop != nullptr);
                     mMetadata->mDiscontinuousLoops.insert(loop);
-                    onDiscontinuousLoop();
                 }
                 break;
               case EOpKill:
@@ -253,7 +259,6 @@ class PullComputeDiscontinuousLoops : public TIntermTraverser
                             mMetadata->mDiscontinuousLoops.insert(loop);
                         }
                     }
-                    onDiscontinuousLoop();
                 }
                 break;
               default:
@@ -274,9 +279,9 @@ class PullComputeDiscontinuousLoops : public TIntermTraverser
                 ASSERT(calleeIndex != CallDAG::InvalidIndex && calleeIndex < mIndex);
                 UNUSED_ASSERTION_VARIABLE(mIndex);
 
-                if ((*mMetadataList)[calleeIndex].mHasDiscontinuousLoopInCallGraph)
+                if ((*mMetadataList)[calleeIndex].mHasGradientLoopInCallGraph)
                 {
-                    onDiscontinuousLoop();
+                    onGradientLoop();
                 }
             }
         }
@@ -375,19 +380,14 @@ class PushDiscontinuousLoops : public TIntermTraverser
 
 }
 
-bool ASTMetadataHLSL::hasGradientInCallGraph(TIntermSelection *node)
-{
-    return mControlFlowsContainingGradient.count(node) > 0;
-}
-
 bool ASTMetadataHLSL::hasGradientInCallGraph(TIntermLoop *node)
 {
     return mControlFlowsContainingGradient.count(node) > 0;
 }
 
-bool ASTMetadataHLSL::hasDiscontinuousLoop(TIntermSelection *node)
+bool ASTMetadataHLSL::hasGradientLoop(TIntermSelection *node)
 {
-    return mIfsContainingDiscontinuousLoop.count(node) > 0;
+    return mIfsContainingGradientLoop.count(node) > 0;
 }
 
 MetadataList CreateASTMetadataHLSL(TIntermNode *root, const CallDAG &callDag)
@@ -424,10 +424,10 @@ MetadataList CreateASTMetadataHLSL(TIntermNode *root, const CallDAG &callDag)
     // of callgraph analysis as for the gradient.
 
     // First compute which loops are discontinuous (no specific order) and pull
-    // the ifs and functions using a discontinuous loop.
+    // the ifs and functions using a gradient loop.
     for (size_t i = 0; i < callDag.size(); i++)
     {
-        PullComputeDiscontinuousLoops pull(&metadataList, i, callDag);
+        PullComputeDiscontinuousAndGradientLoops pull(&metadataList, i, callDag);
         pull.traverse(callDag.getRecordFromIndex(i).node);
     }
 

src/compiler/translator/ASTMetadataHLSL.h

@@ -22,16 +22,15 @@ struct ASTMetadataHLSL
     ASTMetadataHLSL()
         : mUsesGradient(false),
           mCalledInDiscontinuousLoop(false),
-          mHasDiscontinuousLoopInCallGraph(false),
+          mHasGradientLoopInCallGraph(false),
           mNeedsLod0(false)
     {
     }
 
     // Here "something uses a gradient" means here that it either contains a
     // gradient operation, or a call to a function that uses a gradient.
-    bool hasGradientInCallGraph(TIntermSelection *node);
     bool hasGradientInCallGraph(TIntermLoop *node);
-    bool hasDiscontinuousLoop(TIntermSelection *node);
+    bool hasGradientLoop(TIntermSelection *node);
 
     // Does the function use a gradient.
     bool mUsesGradient;
@@ -43,9 +42,9 @@ struct ASTMetadataHLSL
     // Remember information about the discontinuous loops and which functions
     // are called in such loops.
     bool mCalledInDiscontinuousLoop;
-    bool mHasDiscontinuousLoopInCallGraph;
+    bool mHasGradientLoopInCallGraph;
     std::set<TIntermLoop*> mDiscontinuousLoops;
-    std::set<TIntermSelection*> mIfsContainingDiscontinuousLoop;
+    std::set<TIntermSelection *> mIfsContainingGradientLoop;
 
     // Will we need to generate a Lod0 version of the function.
     bool mNeedsLod0;

src/compiler/translator/OutputHLSL.cpp

@@ -2374,9 +2374,7 @@ bool OutputHLSL::visitSelection(Visit visit, TIntermSelection *node)
     }
 
     // D3D errors when there is a gradient operation in a loop in an unflattened if.
-    if (mShaderType == GL_FRAGMENT_SHADER &&
-        mCurrentFunctionMetadata->hasDiscontinuousLoop(node) &&
-        mCurrentFunctionMetadata->hasGradientInCallGraph(node))
+    if (mShaderType == GL_FRAGMENT_SHADER && mCurrentFunctionMetadata->hasGradientLoop(node))
     {
         out << "FLATTEN ";
     }

src/tests/compiler_tests/UnrollFlatten_test.cpp

@@ -105,7 +105,7 @@ TEST_F(UnrollFlattenTest, NoGradient)
     // 2 - no FLATTEN because does not contain discont loop
     // 3 - shouldn't get a Lod0 version generated
     // 4 - no LOOP because discont, and also no gradient
-    // 5 - no FLATTEN because does not contain discont loop
+    // 5 - no FLATTEN because does not contain loop with a gradient
     // 6 - call non-Lod0 version
     // 7 - no FLATTEN
     const char *expectations[] =
@@ -146,16 +146,16 @@ TEST_F(UnrollFlattenTest, GradientNotInDiscont)
     // 2 - no Lod0 version generated
     // 3 - shouldn't get a Lod0 version generated (not in discont loop)
     // 4 - should have LOOP because it contains a gradient operation (even if Lod0)
-    // 5 - no FLATTEN because doesn't contain discont loop
+    // 5 - no FLATTEN because doesn't contain loop with a gradient
     // 6 - call non-Lod0 version
     // 7 - call non-Lod0 version
-    // 8 - no FLATTEN
+    // 8 - FLATTEN because it contains a loop with a gradient
     compile(shaderString);
     const char *expectations[] =
     {
         "fun(", "texture2D(",
         "fun2(", "LOOP", "for", "if", "fun(", "texture2D(",
-        "main(", "if", "fun2("
+        "main(", "FLATTEN", "if", "fun2("
     };
     expect(expectations, ArraySize(expectations));
 }
@@ -188,10 +188,10 @@ TEST_F(UnrollFlattenTest, GradientInDiscont)
     // 2 - will get the Lod0 if in funLod0
     // 3 - shouldn't get a Lod0 version generated (not in discont loop)
     // 4 - should have LOOP because it contains a gradient operation (even if Lod0)
-    // 5 - no FLATTEN because doesn't contain discont loop
+    // 5 - no FLATTEN because doesn't contain a loop with a gradient
     // 6 - call Lod0 version
     // 7 - call Lod0 version
-    // 8 - should have a FLATTEN because has a discont loop and gradient
+    // 8 - FLATTEN because it contains a loop with a gradient
     compile(shaderString);
     const char *expectations[] =
     {