Rewrite repeated assignments to swizzled vectors on NVIDIA

include/GLSLANG/ShaderLang.h

@@ -261,6 +261,10 @@ const ShCompileOptions SH_SKIP_D3D_CONSTANT_REGISTER_ZERO = UINT64_C(1) << 37;
 // Clamp gl_FragDepth to the range [0.0, 1.0] in case it is statically used.
 const ShCompileOptions SH_CLAMP_FRAG_DEPTH = UINT64_C(1) << 38;
 
+// Rewrite expressions like "v.x = z = expression;". Works around a bug in NVIDIA OpenGL drivers
+// prior to version 397.31.
+const ShCompileOptions SH_REWRITE_REPEATED_ASSIGN_TO_SWIZZLED = UINT64_C(1) << 39;
+
 // Defines alternate strategies for implementing array index clamping.
 enum ShArrayIndexClampingStrategy
 {

src/compiler.gypi

@@ -154,6 +154,8 @@
             'compiler/translator/tree_ops/RemoveUnreferencedVariables.h',
             'compiler/translator/tree_ops/RewriteDoWhile.cpp',
             'compiler/translator/tree_ops/RewriteDoWhile.h',
+            'compiler/translator/tree_ops/RewriteRepeatedAssignToSwizzled.cpp',
+            'compiler/translator/tree_ops/RewriteRepeatedAssignToSwizzled.h',
             'compiler/translator/tree_ops/RewriteTexelFetchOffset.cpp',
             'compiler/translator/tree_ops/RewriteTexelFetchOffset.h',
             'compiler/translator/tree_ops/RewriteUnaryMinusOperatorFloat.cpp',

src/compiler/translator/Compiler.cpp

@@ -38,6 +38,7 @@
 #include "compiler/translator/tree_ops/RemovePow.h"
 #include "compiler/translator/tree_ops/RemoveUnreferencedVariables.h"
 #include "compiler/translator/tree_ops/RewriteDoWhile.h"
+#include "compiler/translator/tree_ops/RewriteRepeatedAssignToSwizzled.h"
 #include "compiler/translator/tree_ops/ScalarizeVecAndMatConstructorArgs.h"
 #include "compiler/translator/tree_ops/SeparateDeclarations.h"
 #include "compiler/translator/tree_ops/SimplifyLoopConditions.h"
@@ -737,6 +738,11 @@ bool TCompiler::checkAndSimplifyAST(TIntermBlock *root,
         ClampFragDepth(root, &getSymbolTable());
     }
 
+    if (compileOptions & SH_REWRITE_REPEATED_ASSIGN_TO_SWIZZLED)
+    {
+        sh::RewriteRepeatedAssignToSwizzled(root);
+    }
+
     if (compileOptions & SH_REWRITE_VECTOR_SCALAR_ARITHMETIC)
     {
         VectorizeVectorScalarArithmetic(root, &getSymbolTable());

src/compiler/translator/tree_ops/RewriteRepeatedAssignToSwizzled.cpp

+//
+// Copyright (c) 2018 The ANGLE Project Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+//
+// RewriteRepeatedAssignToSwizzled.cpp: Rewrite expressions that assign an assignment to a swizzled
+// vector, like:
+//     v.x = z = expression;
+// to:
+//     z = expression;
+//     v.x = z;
+//
+// Note that this doesn't handle some corner cases: expressions nested inside other expressions,
+// inside loop headers, or inside if conditions.
+
+#include "compiler/translator/tree_ops/RewriteRepeatedAssignToSwizzled.h"
+
+#include "compiler/translator/tree_util/IntermNode_util.h"
+#include "compiler/translator/tree_util/IntermTraverse.h"
+
+namespace sh
+{
+
+namespace
+{
+
+class RewriteAssignToSwizzledTraverser : public TIntermTraverser
+{
+  public:
+    static void rewrite(TIntermBlock *root);
+
+  private:
+    RewriteAssignToSwizzledTraverser();
+
+    bool visitBinary(Visit, TIntermBinary *node) override;
+
+    void nextIteration();
+
+    bool didRewrite() { return mDidRewrite; }
+
+    bool mDidRewrite;
+};
+
+// static
+void RewriteAssignToSwizzledTraverser::rewrite(TIntermBlock *root)
+{
+    RewriteAssignToSwizzledTraverser rewrite;
+    do
+    {
+        rewrite.nextIteration();
+        root->traverse(&rewrite);
+        rewrite.updateTree();
+    } while (rewrite.didRewrite());
+}
+
+RewriteAssignToSwizzledTraverser::RewriteAssignToSwizzledTraverser()
+    : TIntermTraverser(true, false, false), mDidRewrite(false)
+{
+}
+
+void RewriteAssignToSwizzledTraverser::nextIteration()
+{
+    mDidRewrite = false;
+}
+
+bool RewriteAssignToSwizzledTraverser::visitBinary(Visit, TIntermBinary *node)
+{
+    TIntermBinary *rightBinary = node->getRight()->getAsBinaryNode();
+    TIntermBlock *parentBlock  = getParentNode()->getAsBlock();
+    if (parentBlock && node->isAssignment() && node->getLeft()->getAsSwizzleNode() && rightBinary &&
+        rightBinary->isAssignment())
+    {
+        TIntermSequence replacements;
+        replacements.push_back(rightBinary);
+        TIntermTyped *rightAssignmentTargetCopy = rightBinary->getLeft()->deepCopy();
+        TIntermBinary *lastAssign =
+            new TIntermBinary(EOpAssign, node->getLeft(), rightAssignmentTargetCopy);
+        replacements.push_back(lastAssign);
+        mMultiReplacements.push_back(NodeReplaceWithMultipleEntry(parentBlock, node, replacements));
+        mDidRewrite = true;
+        return false;
+    }
+    return true;
+}
+
+}  // anonymous namespace
+
+void RewriteRepeatedAssignToSwizzled(TIntermBlock *root)
+{
+    RewriteAssignToSwizzledTraverser::rewrite(root);
+}
+
+}  // namespace sh

src/compiler/translator/tree_ops/RewriteRepeatedAssignToSwizzled.h

+//
+// Copyright (c) 2018 The ANGLE Project Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+//
+// RewriteRepeatedAssignToSwizzled.h: Rewrite expressions that assign an assignment to a swizzled
+// vector, like:
+//     v.x = z = expression;
+// to:
+//     z = expression;
+//     v.x = z;
+//
+// Note that this doesn't handle some corner cases: expressions nested inside other expressions,
+// inside loop headers, or inside if conditions.
+
+#ifndef COMPILER_TRANSLATOR_TREEOPS_REWRITEREPEATEDASSIGNTOSWIZZLED_H_
+#define COMPILER_TRANSLATOR_TREEOPS_REWRITEREPEATEDASSIGNTOSWIZZLED_H_
+
+namespace sh
+{
+
+class TIntermBlock;
+
+void RewriteRepeatedAssignToSwizzled(TIntermBlock *root);
+
+}  // namespace sh
+
+#endif  // COMPILER_TRANSLATOR_TREEOPS_REWRITEREPEATEDASSIGNTOSWIZZLED_H_

src/libANGLE/renderer/gl/ShaderGL.cpp

@@ -132,6 +132,11 @@ ShCompileOptions ShaderGL::prepareSourceAndReturnOptions(std::stringstream *sour
         options |= SH_CLAMP_FRAG_DEPTH;
     }
 
+    if (mWorkarounds.rewriteRepeatedAssignToSwizzled)
+    {
+        options |= SH_REWRITE_REPEATED_ASSIGN_TO_SWIZZLED;
+    }
+
     if (mMultiviewImplementationType == MultiviewImplementationTypeGL::NV_VIEWPORT_ARRAY2)
     {
         options |= SH_INITIALIZE_BUILTINS_FOR_INSTANCED_MULTIVIEW;

src/libANGLE/renderer/gl/WorkaroundsGL.h

@@ -147,6 +147,11 @@ struct WorkaroundsGL
     // On some NVIDIA drivers gl_FragDepth is not clamped correctly when rendering to a floating
     // point depth buffer. Clamp it in the translated shader to fix this.
     bool clampFragDepth = false;
+
+    // On some NVIDIA drivers before version 397.31 repeated assignment to swizzled values inside a
+    // GLSL user-defined function have incorrect results. Rewrite this type of statements to fix
+    // this.
+    bool rewriteRepeatedAssignToSwizzled = false;
 };
 
 inline WorkaroundsGL::WorkaroundsGL() = default;

src/libANGLE/renderer/gl/renderergl_utils.cpp

@@ -1161,6 +1161,10 @@ void GenerateWorkarounds(const FunctionsGL *functions, WorkaroundsGL *workaround
     // 390 are known to be affected. Versions after that are expected not to be affected.
     workarounds->clampFragDepth = IsNvidia(vendor);
 
+    // TODO(oetuaho): Make this specific to the affected driver versions. Versions since 397.31 are
+    // not affected.
+    workarounds->rewriteRepeatedAssignToSwizzled = IsNvidia(vendor);
+
 #if defined(ANGLE_PLATFORM_ANDROID)
     // TODO(jmadill): Narrow workaround range for specific devices.
     workarounds->reapplyUBOBindingsAfterUsingBinaryProgram = true;

src/tests/gl_tests/GLSLTest.cpp

@@ -4253,6 +4253,39 @@ TEST_P(GLSLTest_ES31, ExceedCombinedShaderOutputResourcesInVSAndFS)
     ASSERT_GL_NO_ERROR();
 }
 
+// Test that assigning an assignment expression to a swizzled vector field in a user-defined
+// function works correctly.
+TEST_P(GLSLTest_ES3, AssignAssignmentToSwizzled)
+{
+    const std::string &fragmentShader =
+        R"(#version 300 es
+
+        precision highp float;
+        out vec4 my_FragColor;
+
+        uniform float uzero;
+
+        vec3 fun(float s, float v)
+        {
+            vec3 r = vec3(0);
+            if (s < 1.0) {
+                r.x = r.y = r.z = v;
+                return r;
+            }
+            return r;
+        }
+
+        void main()
+        {
+            my_FragColor.a = 1.0;
+            my_FragColor.rgb = fun(uzero, 1.0);
+        })";
+
+    ANGLE_GL_PROGRAM(program, essl3_shaders::vs::Simple(), fragmentShader);
+    drawQuad(program.get(), essl3_shaders::PositionAttrib(), 0.5f);
+    EXPECT_PIXEL_COLOR_EQ(0, 0, GLColor::white);
+}
+
 // Use this to select which configurations (e.g. which renderer, which GLES major version) these
 // tests should be run against.
 ANGLE_INSTANTIATE_TEST(GLSLTest,