ES31: Add atomic_uint support to HLSL translator

This is the first commit in a series to enable atomic counter buffers.

Adds support for atomic counters to the GLSL->HLSL translator using
RWByteAddressBuffer.

Bug: angleproject:1729
Test: angle_end2end_tests

Change-Id: I3b7e08f9256dc9bdbcc02ad8910040f2bc14aeac
Reviewed-on: https://chromium-review.googlesource.com/c/1291329
Commit-Queue: Corentin Wallez <cwallez@chromium.org>
Reviewed-by: Corentin Wallez <cwallez@chromium.org>
diff --git a/src/compiler/translator/AtomicCounterFunctionHLSL.cpp b/src/compiler/translator/AtomicCounterFunctionHLSL.cpp
new file mode 100644
index 0000000..63fe026
--- /dev/null
+++ b/src/compiler/translator/AtomicCounterFunctionHLSL.cpp
@@ -0,0 +1,100 @@
+//
+// 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.
+//
+// AtomicCounterFunctionHLSL: Class for writing implementation of atomic counter functions into HLSL
+// output.
+//
+
+#include "compiler/translator/AtomicCounterFunctionHLSL.h"
+
+#include "compiler/translator/ImmutableStringBuilder.h"
+#include "compiler/translator/InfoSink.h"
+#include "compiler/translator/IntermNode.h"
+
+namespace sh
+{
+
+namespace
+{
+constexpr ImmutableString kAtomicCounter("atomicCounter");
+constexpr ImmutableString kAtomicCounterIncrement("atomicCounterIncrement");
+constexpr ImmutableString kAtomicCounterDecrement("atomicCounterDecrement");
+constexpr ImmutableString kAtomicCounterBaseName("_acbase_");
+}  // namespace
+
+ImmutableString AtomicCounterFunctionHLSL::useAtomicCounterFunction(const ImmutableString &name)
+{
+    // The largest string that will be create created is "_acbase_increment" or "_acbase_decrement"
+    ImmutableStringBuilder hlslFunctionNameSB(kAtomicCounterBaseName.length() +
+                                              strlen("increment"));
+    hlslFunctionNameSB << kAtomicCounterBaseName;
+
+    AtomicCounterFunction atomicMethod;
+    if (kAtomicCounter == name)
+    {
+        atomicMethod = AtomicCounterFunction::LOAD;
+        hlslFunctionNameSB << "load";
+    }
+    else if (kAtomicCounterIncrement == name)
+    {
+        atomicMethod = AtomicCounterFunction::INCREMENT;
+        hlslFunctionNameSB << "increment";
+    }
+    else if (kAtomicCounterDecrement == name)
+    {
+        atomicMethod = AtomicCounterFunction::DECREMENT;
+        hlslFunctionNameSB << "decrement";
+    }
+    else
+    {
+        atomicMethod = AtomicCounterFunction::INVALID;
+        UNREACHABLE();
+    }
+
+    ImmutableString hlslFunctionName(hlslFunctionNameSB);
+    mAtomicCounterFunctions[hlslFunctionName] = atomicMethod;
+
+    return hlslFunctionName;
+}
+
+void AtomicCounterFunctionHLSL::atomicCounterFunctionHeader(TInfoSinkBase &out)
+{
+    for (auto &atomicFunction : mAtomicCounterFunctions)
+    {
+        out << "uint " << atomicFunction.first
+            << "(in RWByteAddressBuffer counter, int address)\n"
+               "{\n";
+        switch (atomicFunction.second)
+        {
+            case AtomicCounterFunction::INCREMENT:
+            case AtomicCounterFunction::DECREMENT:
+                out << "    uint ret;\n"
+                       "    counter.InterlockedAdd(address, ";
+                if (atomicFunction.second == AtomicCounterFunction::DECREMENT)
+                {
+                    out << "0u - ";
+                }
+                out << "1u, ret);\n"
+                    << "    return ret;\n";
+                break;
+            case AtomicCounterFunction::LOAD:
+                out << "    return counter.Load(address);\n";
+                break;
+            default:
+                UNREACHABLE();
+                break;
+        }
+        out << "}\n\n";
+    }
+}
+
+ImmutableString getAtomicCounterNameForBinding(int binding)
+{
+    std::stringstream counterName;
+    counterName << kAtomicCounterBaseName << binding;
+    return ImmutableString(counterName.str());
+}
+
+}  // namespace sh
diff --git a/src/compiler/translator/AtomicCounterFunctionHLSL.h b/src/compiler/translator/AtomicCounterFunctionHLSL.h
new file mode 100644
index 0000000..63fd891
--- /dev/null
+++ b/src/compiler/translator/AtomicCounterFunctionHLSL.h
@@ -0,0 +1,47 @@
+//
+// 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.
+//
+// AtomicCounterFunctionHLSL: Class for writing implementation of atomic counter functions into HLSL
+// output.
+//
+
+#ifndef COMPILER_TRANSLATOR_ATOMICCOUNTERFUNCTIONHLSL_H_
+#define COMPILER_TRANSLATOR_ATOMICCOUNTERFUNCTIONHLSL_H_
+
+#include <map>
+
+#include "compiler/translator/Common.h"
+#include "compiler/translator/ImmutableString.h"
+
+namespace sh
+{
+
+class TInfoSinkBase;
+struct TLayoutQualifier;
+
+class AtomicCounterFunctionHLSL final : angle::NonCopyable
+{
+  public:
+    ImmutableString useAtomicCounterFunction(const ImmutableString &name);
+
+    void atomicCounterFunctionHeader(TInfoSinkBase &out);
+
+  private:
+    enum class AtomicCounterFunction
+    {
+        LOAD,
+        INCREMENT,
+        DECREMENT,
+        INVALID
+    };
+
+    std::map<ImmutableString, AtomicCounterFunction> mAtomicCounterFunctions;
+};
+
+ImmutableString getAtomicCounterNameForBinding(int binding);
+
+}  // namespace sh
+
+#endif  // COMPILER_TRANSLATOR_ATOMICCOUNTERFUNCTIONHLSL_H_
diff --git a/src/compiler/translator/OutputHLSL.cpp b/src/compiler/translator/OutputHLSL.cpp
index 457f412..fee56a4 100644
--- a/src/compiler/translator/OutputHLSL.cpp
+++ b/src/compiler/translator/OutputHLSL.cpp
@@ -13,6 +13,7 @@
 #include "common/angleutils.h"
 #include "common/debug.h"
 #include "common/utilities.h"
+#include "compiler/translator/AtomicCounterFunctionHLSL.h"
 #include "compiler/translator/BuiltInFunctionEmulator.h"
 #include "compiler/translator/BuiltInFunctionEmulatorHLSL.h"
 #include "compiler/translator/ImageFunctionHLSL.h"
@@ -233,9 +234,10 @@
 
     mExcessiveLoopIndex = nullptr;
 
-    mStructureHLSL       = new StructureHLSL;
-    mTextureFunctionHLSL = new TextureFunctionHLSL;
-    mImageFunctionHLSL   = new ImageFunctionHLSL;
+    mStructureHLSL             = new StructureHLSL;
+    mTextureFunctionHLSL       = new TextureFunctionHLSL;
+    mImageFunctionHLSL         = new ImageFunctionHLSL;
+    mAtomicCounterFunctionHLSL = new AtomicCounterFunctionHLSL;
 
     unsigned int firstUniformRegister =
         ((compileOptions & SH_SKIP_D3D_CONSTANT_REGISTER_ZERO) != 0) ? 1u : 0u;
@@ -263,6 +265,7 @@
     SafeDelete(mResourcesHLSL);
     SafeDelete(mTextureFunctionHLSL);
     SafeDelete(mImageFunctionHLSL);
+    SafeDelete(mAtomicCounterFunctionHLSL);
     for (auto &eqFunction : mStructEqualityFunctions)
     {
         SafeDelete(eqFunction);
@@ -545,6 +548,11 @@
            "#define FLATTEN\n"
            "#endif\n";
 
+    // array stride for atomic counter buffers is always 4 per original extension
+    // ARB_shader_atomic_counters and discussion on
+    // https://github.com/KhronosGroup/OpenGL-API/issues/5
+    out << "\n#define ATOMIC_COUNTER_ARRAY_STRIDE 4\n\n";
+
     if (mShaderType == GL_FRAGMENT_SHADER)
     {
         const bool usingMRTExtension =
@@ -864,6 +872,7 @@
         (mCompileOptions & SH_HLSL_GET_DIMENSIONS_IGNORES_BASE_LEVEL) != 0;
     mTextureFunctionHLSL->textureFunctionHeader(out, mOutputType, getDimensionsIgnoresBaseLevel);
     mImageFunctionHLSL->imageFunctionHeader(out);
+    mAtomicCounterFunctionHLSL->atomicCounterFunctionHeader(out);
 
     if (mUsesFragCoord)
     {
@@ -941,6 +950,21 @@
         mUsesDepthRange = true;
         out << name;
     }
+    else if (IsAtomicCounter(variable.getType().getBasicType()))
+    {
+        const TType &variableType = variable.getType();
+        if (variableType.getQualifier() == EvqUniform)
+        {
+            TLayoutQualifier layout             = variableType.getLayoutQualifier();
+            mReferencedUniforms[uniqueId.get()] = &variable;
+            out << getAtomicCounterNameForBinding(layout.binding) << ", " << layout.offset;
+        }
+        else
+        {
+            TString varName = DecorateVariableIfNeeded(variable);
+            out << varName << ", " << varName << "_offset";
+        }
+    }
     else
     {
         const TType &variableType = variable.getType();
@@ -1377,6 +1401,10 @@
                 // separator to access the sampler variable that has been moved out of the struct.
                 outputTriplet(out, visit, "", "_", "");
             }
+            else if (IsAtomicCounter(leftType.getBasicType()))
+            {
+                outputTriplet(out, visit, "", " + (", ") * ATOMIC_COUNTER_ARRAY_STRIDE");
+            }
             else
             {
                 outputTriplet(out, visit, "", "[", "]");
@@ -1384,10 +1412,21 @@
         }
         break;
         case EOpIndexIndirect:
+        {
             // We do not currently support indirect references to interface blocks
             ASSERT(node->getLeft()->getBasicType() != EbtInterfaceBlock);
-            outputTriplet(out, visit, "", "[", "]");
+
+            const TType &leftType = node->getLeft()->getType();
+            if (IsAtomicCounter(leftType.getBasicType()))
+            {
+                outputTriplet(out, visit, "", " + (", ") * ATOMIC_COUNTER_ARRAY_STRIDE");
+            }
+            else
+            {
+                outputTriplet(out, visit, "", "[", "]");
+            }
             break;
+        }
         case EOpIndexDirectStruct:
         {
             const TStructure *structure       = node->getLeft()->getType().getStruct();
@@ -2111,6 +2150,13 @@
                     type.getMemoryQualifier().readonly);
                 out << imageFunctionName << "(";
             }
+            else if (node->getFunction()->isAtomicCounterFunction())
+            {
+                const ImmutableString &name = node->getFunction()->name();
+                ImmutableString atomicFunctionName =
+                    mAtomicCounterFunctionHLSL->useAtomicCounterFunction(name);
+                out << atomicFunctionName << "(";
+            }
             else
             {
                 const ImmutableString &name = node->getFunction()->name();
@@ -2871,8 +2917,18 @@
         }
     }
 
-    out << QualifierString(qualifier) << " " << TypeString(type) << " " << nameStr
-        << ArrayString(type);
+    // If the parameter is an atomic counter, we need to add an extra parameter to keep track of the
+    // buffer offset.
+    if (IsAtomicCounter(type.getBasicType()))
+    {
+        out << QualifierString(qualifier) << " " << TypeString(type) << " " << nameStr << ", int "
+            << nameStr << "_offset";
+    }
+    else
+    {
+        out << QualifierString(qualifier) << " " << TypeString(type) << " " << nameStr
+            << ArrayString(type);
+    }
 
     // If the structure parameter contains samplers, they need to be passed into the function as
     // separate parameters. HLSL doesn't natively support samplers in structs.
diff --git a/src/compiler/translator/OutputHLSL.h b/src/compiler/translator/OutputHLSL.h
index b861cca..00608b6 100644
--- a/src/compiler/translator/OutputHLSL.h
+++ b/src/compiler/translator/OutputHLSL.h
@@ -23,6 +23,7 @@
 
 namespace sh
 {
+class AtomicCounterFunctionHLSL;
 class ImageFunctionHLSL;
 class ResourcesHLSL;
 class StructureHLSL;
@@ -177,6 +178,7 @@
     ResourcesHLSL *mResourcesHLSL;
     TextureFunctionHLSL *mTextureFunctionHLSL;
     ImageFunctionHLSL *mImageFunctionHLSL;
+    AtomicCounterFunctionHLSL *mAtomicCounterFunctionHLSL;
 
     // Parameters determining what goes in the header output
     bool mUsesFragColor;
diff --git a/src/compiler/translator/ResourcesHLSL.cpp b/src/compiler/translator/ResourcesHLSL.cpp
index 07223cf..0964f33 100644
--- a/src/compiler/translator/ResourcesHLSL.cpp
+++ b/src/compiler/translator/ResourcesHLSL.cpp
@@ -10,6 +10,7 @@
 #include "compiler/translator/ResourcesHLSL.h"
 
 #include "common/utilities.h"
+#include "compiler/translator/AtomicCounterFunctionHLSL.h"
 #include "compiler/translator/ImmutableStringBuilder.h"
 #include "compiler/translator/StructureHLSL.h"
 #include "compiler/translator/UtilsHLSL.h"
@@ -378,6 +379,14 @@
     out << ArrayString(type) << " : " << registerString << ";\n";
 }
 
+void ResourcesHLSL::outputAtomicCounterBuffer(TInfoSinkBase &out,
+                                              const int binding,
+                                              const unsigned int registerIndex)
+{
+    out << "uniform RWByteAddressBuffer " << getAtomicCounterNameForBinding(binding)
+        << " : register(u" << registerIndex << ");\n";
+}
+
 void ResourcesHLSL::uniformsHeader(TInfoSinkBase &out,
                                    ShShaderOutput outputType,
                                    const ReferencedVariables &referencedUniforms,
@@ -394,6 +403,8 @@
     TMap<const TVariable *, TString> samplerInStructSymbolsToAPINames;
     TVector<TVector<const TVariable *>> groupedReadonlyImageUniforms(HLSL_TEXTURE_MAX + 1);
     TVector<TVector<const TVariable *>> groupedImageUniforms(HLSL_RWTEXTURE_MAX + 1);
+
+    TUnorderedMap<int, unsigned int> assignedAtomicCounterBindings;
     for (auto &uniformIt : referencedUniforms)
     {
         // Output regular uniforms. Group sampler uniforms by type.
@@ -425,6 +436,24 @@
                 groupedImageUniforms[group].push_back(&variable);
             }
         }
+        else if (outputType == SH_HLSL_4_1_OUTPUT && IsAtomicCounter(type.getBasicType()))
+        {
+            TLayoutQualifier layout = type.getLayoutQualifier();
+            int binding             = layout.binding;
+            unsigned int registerIndex;
+            if (assignedAtomicCounterBindings.find(binding) == assignedAtomicCounterBindings.end())
+            {
+                registerIndex                          = mUAVRegister++;
+                assignedAtomicCounterBindings[binding] = registerIndex;
+                outputAtomicCounterBuffer(out, binding, registerIndex);
+            }
+            else
+            {
+                registerIndex = assignedAtomicCounterBindings[binding];
+            }
+            const Uniform *uniform             = findUniformByName(variable.name());
+            mUniformRegisterMap[uniform->name] = registerIndex;
+        }
         else
         {
             if (type.isStructureContainingSamplers())
@@ -469,8 +498,10 @@
 
     if (outputType == SH_HLSL_4_1_OUTPUT)
     {
-        unsigned int groupTextureRegisterIndex   = 0;
-        unsigned int groupRWTextureRegisterIndex = 0;
+        unsigned int groupTextureRegisterIndex = 0;
+        // Atomic counters and RW texture share the same resources. Therefore, RW texture need to
+        // start counting after the last atomic counter.
+        unsigned int groupRWTextureRegisterIndex = mUAVRegister;
         unsigned int imageUniformGroupIndex      = 0;
         // TEXTURE_2D is special, index offset is assumed to be 0 and omitted in that case.
         ASSERT(HLSL_TEXTURE_MIN == HLSL_TEXTURE_2D);
@@ -693,4 +724,4 @@
            "{\n" +
            uniformBlockMembersString(interfaceBlock, blockStorage) + "};\n\n";
 }
-}
+}  // namespace sh
diff --git a/src/compiler/translator/ResourcesHLSL.h b/src/compiler/translator/ResourcesHLSL.h
index 5dd02b2..567ae6d 100644
--- a/src/compiler/translator/ResourcesHLSL.h
+++ b/src/compiler/translator/ResourcesHLSL.h
@@ -81,6 +81,9 @@
                        const TType &type,
                        const TVariable &variable,
                        const unsigned int registerIndex);
+    void outputAtomicCounterBuffer(TInfoSinkBase &out,
+                                   const int binding,
+                                   const unsigned int registerIndex);
 
     // Returns the uniform's register index
     unsigned int assignUniformRegister(const TType &type,
diff --git a/src/compiler/translator/Symbol.cpp b/src/compiler/translator/Symbol.cpp
index 6923568..20754f5 100644
--- a/src/compiler/translator/Symbol.cpp
+++ b/src/compiler/translator/Symbol.cpp
@@ -25,6 +25,7 @@
 constexpr const ImmutableString kImageLoadName("imageLoad");
 constexpr const ImmutableString kImageStoreName("imageStore");
 constexpr const ImmutableString kImageSizeName("imageSize");
+constexpr const ImmutableString kAtomicCounterName("atomicCounter");
 
 static const char kFunctionMangledNameSeparator = '(';
 
@@ -219,6 +220,11 @@
            (name() == kImageSizeName || name() == kImageLoadName || name() == kImageStoreName);
 }
 
+bool TFunction::isAtomicCounterFunction() const
+{
+    return SymbolType() == SymbolType::BuiltIn && name().beginsWith(kAtomicCounterName);
+}
+
 bool TFunction::hasSamplerInStructParams() const
 {
     for (size_t paramIndex = 0; paramIndex < mParamCount; ++paramIndex)
diff --git a/src/compiler/translator/Symbol.h b/src/compiler/translator/Symbol.h
index d839b12..7ae5b97 100644
--- a/src/compiler/translator/Symbol.h
+++ b/src/compiler/translator/Symbol.h
@@ -235,6 +235,7 @@
 
     bool isMain() const;
     bool isImageFunction() const;
+    bool isAtomicCounterFunction() const;
     bool hasSamplerInStructParams() const;
 
     // Note: Only to be used for static built-in functions!
diff --git a/src/compiler/translator/UtilsHLSL.cpp b/src/compiler/translator/UtilsHLSL.cpp
index 5e1c497..53d471e 100644
--- a/src/compiler/translator/UtilsHLSL.cpp
+++ b/src/compiler/translator/UtilsHLSL.cpp
@@ -967,7 +967,8 @@
             case EbtSamplerExternalOES:
                 return "sampler2D";
             case EbtAtomicCounter:
-                return "atomic_uint";
+                // Multiple atomic_uints will be implemented as a single RWByteAddressBuffer
+                return "RWByteAddressBuffer";
             default:
                 break;
         }