ES31: Use indices to access image variables in built-in image functions

In order to implement glBindImageTexture to bind a layer of 3D/2DArray/Cube
texture, use indices to access image variables when translating built-in
image functions.
There is a conflict when transferring image2D/iimage2D/uimage2D variables to
an user defined function. For example,
layout(r32ui, binding = 0) readonly uniform highp uimage2D uImage_1;
layout(r32ui, binding = 1) readonly uniform highp uimage2D uImage_2;
uvec4 lod_fun(uimage2D img, ivec2 p)
{
    return imageLoad(img, p);
}
void main()
{
    uvec4 value_1 = lod_fun(uImage_1, ivec2(gl_LocalInvocationID.xy));
    uvec4 value_2 = lod_fun(uImage_2, ivec2(gl_LocalInvocationID.xy));
}
If uImage_1 binds to a 2D texture, and uImage_2 binds to a layer of 3D texture,
uImage_1 will be translated to Texture2D type, and uImage_2 will be translated to
Texture3D type, "img" type of lod_fun will be translated Texture2D, so uImage_2
cannot be transferred to lod_fun as a parameter.
Indices without Texture/RWTexture information could handle this situation easily.

BUG=angleproject:1987
TEST=angle_end2end_tests.ComputeShaderTest.*

Change-Id: I7647395f0042f613c5d6e9eeb49392ab6252e21e
Reviewed-on: https://chromium-review.googlesource.com/1065797
Commit-Queue: Xinghua Cao <xinghua.cao@intel.com>
Reviewed-by: Geoff Lang <geofflang@chromium.org>
Reviewed-by: Olli Etuaho <oetuaho@nvidia.com>
diff --git a/src/compiler/translator/ImageFunctionHLSL.cpp b/src/compiler/translator/ImageFunctionHLSL.cpp
index b3c3931..64e4c0a 100644
--- a/src/compiler/translator/ImageFunctionHLSL.cpp
+++ b/src/compiler/translator/ImageFunctionHLSL.cpp
@@ -7,24 +7,48 @@
 //
 
 #include "compiler/translator/ImageFunctionHLSL.h"
+#include "compiler/translator/ImmutableStringBuilder.h"
 #include "compiler/translator/UtilsHLSL.h"
 
 namespace sh
 {
 
 // static
+ImmutableString ImageFunctionHLSL::GetImageReference(
+    TInfoSinkBase &out,
+    const ImageFunctionHLSL::ImageFunction &imageFunction)
+{
+    static const ImmutableString kImageIndexStr("[index]");
+    if (imageFunction.readonly)
+    {
+        static const ImmutableString kReadonlyImagesStr("readonlyImages");
+        ImmutableString suffix(
+            TextureGroupSuffix(imageFunction.image, imageFunction.imageInternalFormat));
+        out << "    const uint index = imageIndex - readonlyImageIndexOffset" << suffix.data()
+            << ";\n";
+        ImmutableStringBuilder imageRefBuilder(kReadonlyImagesStr.length() + suffix.length() +
+                                               kImageIndexStr.length());
+        imageRefBuilder << kReadonlyImagesStr << suffix << kImageIndexStr;
+        return imageRefBuilder;
+    }
+    else
+    {
+        static const ImmutableString kImagesStr("images");
+        ImmutableString suffix(
+            RWTextureGroupSuffix(imageFunction.image, imageFunction.imageInternalFormat));
+        out << "    const uint index = imageIndex - imageIndexOffset" << suffix.data() << ";\n";
+        ImmutableStringBuilder imageRefBuilder(kImagesStr.length() + suffix.length() +
+                                               kImageIndexStr.length());
+        imageRefBuilder << kImagesStr << suffix << kImageIndexStr;
+        return imageRefBuilder;
+    }
+}
+
 void ImageFunctionHLSL::OutputImageFunctionArgumentList(
     TInfoSinkBase &out,
     const ImageFunctionHLSL::ImageFunction &imageFunction)
 {
-    if (imageFunction.readonly)
-    {
-        out << TextureString(imageFunction.image, imageFunction.imageInternalFormat) << " tex";
-    }
-    else
-    {
-        out << RWTextureString(imageFunction.image, imageFunction.imageInternalFormat) << " tex";
-    }
+    out << "uint imageIndex";
 
     if (imageFunction.method == ImageFunctionHLSL::ImageFunction::Method::LOAD ||
         imageFunction.method == ImageFunctionHLSL::ImageFunction::Method::STORE)
@@ -84,7 +108,7 @@
 void ImageFunctionHLSL::OutputImageSizeFunctionBody(
     TInfoSinkBase &out,
     const ImageFunctionHLSL::ImageFunction &imageFunction,
-    const TString &imageReference)
+    const ImmutableString &imageReference)
 {
     if (IsImage3D(imageFunction.image) || IsImage2DArray(imageFunction.image) ||
         IsImageCube(imageFunction.image))
@@ -115,7 +139,7 @@
 void ImageFunctionHLSL::OutputImageLoadFunctionBody(
     TInfoSinkBase &out,
     const ImageFunctionHLSL::ImageFunction &imageFunction,
-    const TString &imageReference)
+    const ImmutableString &imageReference)
 {
     if (IsImage3D(imageFunction.image) || IsImage2DArray(imageFunction.image) ||
         IsImageCube(imageFunction.image))
@@ -134,7 +158,7 @@
 void ImageFunctionHLSL::OutputImageStoreFunctionBody(
     TInfoSinkBase &out,
     const ImageFunctionHLSL::ImageFunction &imageFunction,
-    const TString &imageReference)
+    const ImmutableString &imageReference)
 {
     if (IsImage3D(imageFunction.image) || IsImage2DArray(imageFunction.image) ||
         IsImage2D(imageFunction.image) || IsImageCube(imageFunction.image))
@@ -175,6 +199,36 @@
     return name;
 }
 
+ImageFunctionHLSL::ImageFunction::DataType ImageFunctionHLSL::ImageFunction::getDataType(
+    TLayoutImageInternalFormat format) const
+{
+    switch (format)
+    {
+        case EiifRGBA32F:
+        case EiifRGBA16F:
+        case EiifR32F:
+            return ImageFunction::DataType::FLOAT4;
+        case EiifRGBA32UI:
+        case EiifRGBA16UI:
+        case EiifRGBA8UI:
+        case EiifR32UI:
+            return ImageFunction::DataType::UINT4;
+        case EiifRGBA32I:
+        case EiifRGBA16I:
+        case EiifRGBA8I:
+        case EiifR32I:
+            return ImageFunction::DataType::INT4;
+        case EiifRGBA8:
+            return ImageFunction::DataType::UNORM_FLOAT4;
+        case EiifRGBA8_SNORM:
+            return ImageFunction::DataType::SNORM_FLOAT4;
+        default:
+            UNREACHABLE();
+    }
+
+    return ImageFunction::DataType::NONE;
+}
+
 const char *ImageFunctionHLSL::ImageFunction::getReturnType() const
 {
     if (method == ImageFunction::Method::SIZE)
@@ -235,8 +289,8 @@
 
 bool ImageFunctionHLSL::ImageFunction::operator<(const ImageFunction &rhs) const
 {
-    return std::tie(image, imageInternalFormat, readonly, method) <
-           std::tie(rhs.image, rhs.imageInternalFormat, rhs.readonly, rhs.method);
+    return std::tie(image, type, method, readonly) <
+           std::tie(rhs.image, rhs.type, rhs.method, rhs.readonly);
 }
 
 TString ImageFunctionHLSL::useImageFunction(const ImmutableString &name,
@@ -249,6 +303,7 @@
     imageFunction.image               = type;
     imageFunction.imageInternalFormat = imageInternalFormat;
     imageFunction.readonly            = readonly;
+    imageFunction.type                = imageFunction.getDataType(imageInternalFormat);
 
     if (name == "imageSize")
     {
@@ -281,8 +336,7 @@
         out << ")\n"
                "{\n";
 
-        TString imageReference("tex");
-
+        ImmutableString imageReference = GetImageReference(out, imageFunction);
         if (imageFunction.method == ImageFunction::Method::SIZE)
         {
             OutputImageSizeFunctionBody(out, imageFunction, imageReference);
diff --git a/src/compiler/translator/ImageFunctionHLSL.h b/src/compiler/translator/ImageFunctionHLSL.h
index 4ff0110..c58213a 100644
--- a/src/compiler/translator/ImageFunctionHLSL.h
+++ b/src/compiler/translator/ImageFunctionHLSL.h
@@ -43,30 +43,45 @@
             STORE
         };
 
+        enum class DataType
+        {
+            NONE,
+            FLOAT4,
+            UINT4,
+            INT4,
+            UNORM_FLOAT4,
+            SNORM_FLOAT4
+        };
+
         TString name() const;
 
         bool operator<(const ImageFunction &rhs) const;
 
+        DataType getDataType(TLayoutImageInternalFormat format) const;
+
         const char *getReturnType() const;
 
         TBasicType image;
         TLayoutImageInternalFormat imageInternalFormat;
         bool readonly;
         Method method;
+        DataType type;
     };
 
+    static ImmutableString GetImageReference(TInfoSinkBase &out,
+                                             const ImageFunctionHLSL::ImageFunction &imageFunction);
     static void OutputImageFunctionArgumentList(
         TInfoSinkBase &out,
         const ImageFunctionHLSL::ImageFunction &imageFunction);
     static void OutputImageSizeFunctionBody(TInfoSinkBase &out,
                                             const ImageFunctionHLSL::ImageFunction &imageFunction,
-                                            const TString &imageReference);
+                                            const ImmutableString &imageReference);
     static void OutputImageLoadFunctionBody(TInfoSinkBase &out,
                                             const ImageFunctionHLSL::ImageFunction &imageFunction,
-                                            const TString &imageReference);
+                                            const ImmutableString &imageReference);
     static void OutputImageStoreFunctionBody(TInfoSinkBase &out,
                                              const ImageFunctionHLSL::ImageFunction &imageFunction,
-                                             const TString &imageReference);
+                                             const ImmutableString &imageReference);
     using ImageFunctionSet = std::set<ImageFunction>;
     ImageFunctionSet mUsesImage;
 };
diff --git a/src/compiler/translator/OutputHLSL.cpp b/src/compiler/translator/OutputHLSL.cpp
index 1f2b066..1c7e090 100644
--- a/src/compiler/translator/OutputHLSL.cpp
+++ b/src/compiler/translator/OutputHLSL.cpp
@@ -202,8 +202,7 @@
 
     unsigned int firstUniformRegister =
         ((compileOptions & SH_SKIP_D3D_CONSTANT_REGISTER_ZERO) != 0) ? 1u : 0u;
-    mUniformHLSL =
-        new UniformHLSL(shaderType, mStructureHLSL, outputType, uniforms, firstUniformRegister);
+    mUniformHLSL = new UniformHLSL(mStructureHLSL, outputType, uniforms, firstUniformRegister);
 
     if (mOutputType == SH_HLSL_3_0_OUTPUT)
     {
diff --git a/src/compiler/translator/UniformHLSL.cpp b/src/compiler/translator/UniformHLSL.cpp
index 33a5d88..37a09e7 100644
--- a/src/compiler/translator/UniformHLSL.cpp
+++ b/src/compiler/translator/UniformHLSL.cpp
@@ -67,7 +67,7 @@
     return DecoratePrivate(interfaceBlock.name()) + "_type";
 }
 
-void OutputSamplerIndexArrayInitializer(TInfoSinkBase &out,
+void OutputUniformIndexArrayInitializer(TInfoSinkBase &out,
                                         const TType &type,
                                         unsigned int startIndex)
 {
@@ -82,7 +82,7 @@
         }
         if (elementType.isArray())
         {
-            OutputSamplerIndexArrayInitializer(out, elementType,
+            OutputUniformIndexArrayInitializer(out, elementType,
                                                startIndex + i * elementType.getArraySizeProduct());
         }
         else
@@ -95,8 +95,7 @@
 
 }  // anonymous namespace
 
-UniformHLSL::UniformHLSL(sh::GLenum shaderType,
-                         StructureHLSL *structureHLSL,
+UniformHLSL::UniformHLSL(StructureHLSL *structureHLSL,
                          ShShaderOutput outputType,
                          const std::vector<Uniform> &uniforms,
                          unsigned int firstUniformRegister)
@@ -105,7 +104,6 @@
       mTextureRegister(0),
       mRWTextureRegister(0),
       mSamplerCount(0),
-      mShaderType(shaderType),
       mStructureHLSL(structureHLSL),
       mOutputType(outputType),
       mUniforms(uniforms)
@@ -236,7 +234,7 @@
         {
             out << "static const uint " << DecorateVariableIfNeeded(*uniform) << ArrayString(type)
                 << " = ";
-            OutputSamplerIndexArrayInitializer(out, type, samplerArrayIndex);
+            OutputUniformIndexArrayInitializer(out, type, samplerArrayIndex);
             out << ";\n";
         }
         else
@@ -263,6 +261,85 @@
     *groupTextureRegisterIndex += groupRegisterCount;
 }
 
+void UniformHLSL::outputHLSLImageUniformIndices(TInfoSinkBase &out,
+                                                const TVector<const TVariable *> &group,
+                                                unsigned int imageArrayIndex,
+                                                unsigned int *groupRegisterCount)
+{
+    for (const TVariable *uniform : group)
+    {
+        const TType &type           = uniform->getType();
+        const ImmutableString &name = uniform->name();
+        unsigned int registerCount  = 0;
+
+        assignUniformRegister(type, name, &registerCount);
+        *groupRegisterCount += registerCount;
+
+        if (type.isArray())
+        {
+            out << "static const uint " << DecorateVariableIfNeeded(*uniform) << ArrayString(type)
+                << " = ";
+            OutputUniformIndexArrayInitializer(out, type, imageArrayIndex);
+            out << ";\n";
+        }
+        else
+        {
+            out << "static const uint " << DecorateVariableIfNeeded(*uniform) << " = "
+                << imageArrayIndex << ";\n";
+        }
+
+        imageArrayIndex += registerCount;
+    }
+}
+
+void UniformHLSL::outputHLSLReadonlyImageUniformGroup(TInfoSinkBase &out,
+                                                      const HLSLTextureGroup textureGroup,
+                                                      const TVector<const TVariable *> &group,
+                                                      unsigned int *groupTextureRegisterIndex,
+                                                      unsigned int *imageUniformGroupIndex)
+{
+    if (group.empty())
+    {
+        return;
+    }
+
+    unsigned int groupRegisterCount = 0;
+    outputHLSLImageUniformIndices(out, group, *imageUniformGroupIndex, &groupRegisterCount);
+
+    TString suffix = TextureGroupSuffix(textureGroup);
+    out << "static const uint readonlyImageIndexOffset" << suffix << " = "
+        << (*imageUniformGroupIndex) << ";\n";
+    out << "uniform " << TextureString(textureGroup) << " readonlyImages" << suffix << "["
+        << groupRegisterCount << "]"
+        << " : register(t" << (*groupTextureRegisterIndex) << ");\n";
+    *groupTextureRegisterIndex += groupRegisterCount;
+    *imageUniformGroupIndex += groupRegisterCount;
+}
+
+void UniformHLSL::outputHLSLImageUniformGroup(TInfoSinkBase &out,
+                                              const HLSLRWTextureGroup textureGroup,
+                                              const TVector<const TVariable *> &group,
+                                              unsigned int *groupTextureRegisterIndex,
+                                              unsigned int *imageUniformGroupIndex)
+{
+    if (group.empty())
+    {
+        return;
+    }
+
+    unsigned int groupRegisterCount = 0;
+    outputHLSLImageUniformIndices(out, group, *imageUniformGroupIndex, &groupRegisterCount);
+
+    TString suffix = RWTextureGroupSuffix(textureGroup);
+    out << "static const uint imageIndexOffset" << suffix << " = " << (*imageUniformGroupIndex)
+        << ";\n";
+    out << "uniform " << RWTextureString(textureGroup) << " images" << suffix << "["
+        << groupRegisterCount << "]"
+        << " : register(u" << (*groupTextureRegisterIndex) << ");\n";
+    *groupTextureRegisterIndex += groupRegisterCount;
+    *imageUniformGroupIndex += groupRegisterCount;
+}
+
 void UniformHLSL::outputHLSL4_0_FL9_3Sampler(TInfoSinkBase &out,
                                              const TType &type,
                                              const TVariable &variable,
@@ -276,46 +353,6 @@
         << str(registerIndex) << ");\n";
 }
 
-void UniformHLSL::outputHLSL4_1_FL11Texture(TInfoSinkBase &out,
-                                            const TType &type,
-                                            const TVariable &variable,
-                                            const unsigned int registerIndex)
-{
-    // TODO(xinghua.cao@intel.com): if image2D variable is bound on one layer of Texture3D or
-    // Texture2DArray. Translate this variable to HLSL Texture3D object or HLSL Texture2DArray
-    // object, or create a temporary Texture2D to save content of the layer and bind the
-    // temporary Texture2D to image2D variable.
-    out << "uniform "
-        << TextureString(type.getBasicType(), type.getLayoutQualifier().imageInternalFormat) << " "
-        << DecorateVariableIfNeeded(variable) << ArrayString(type) << " : register(t"
-        << str(registerIndex) << ");\n";
-    return;
-}
-
-void UniformHLSL::outputHLSL4_1_FL11RWTexture(TInfoSinkBase &out,
-                                              const TType &type,
-                                              const TVariable &variable,
-                                              const unsigned int registerIndex)
-{
-    // TODO(xinghua.cao@intel.com): if image2D variable is bound on one layer of Texture3D or
-    // Texture2DArray. Translate this variable to HLSL RWTexture3D object or HLSL RWTexture2DArray
-    // object, or create a temporary Texture2D to save content of the layer and bind the
-    // temporary Texture2D to image2D variable.
-    if (mShaderType == GL_COMPUTE_SHADER)
-    {
-        out << "uniform "
-            << RWTextureString(type.getBasicType(), type.getLayoutQualifier().imageInternalFormat)
-            << " " << DecorateVariableIfNeeded(variable) << ArrayString(type) << " : register(u"
-            << str(registerIndex) << ");\n";
-    }
-    else
-    {
-        // TODO(xinghua.cao@intel.com): Support images in vertex shader and fragment shader,
-        // which are needed to sync binding value when linking program.
-    }
-    return;
-}
-
 void UniformHLSL::outputUniform(TInfoSinkBase &out,
                                 const TType &type,
                                 const TVariable &variable,
@@ -355,7 +392,8 @@
     // HLSL sampler type, enumerated in HLSLTextureSamplerGroup.
     TVector<TVector<const TVariable *>> groupedSamplerUniforms(HLSL_TEXTURE_MAX + 1);
     TMap<const TVariable *, TString> samplerInStructSymbolsToAPINames;
-    TVector<const TVariable *> imageUniformsHLSL41Output;
+    TVector<TVector<const TVariable *>> groupedReadonlyImageUniforms(HLSL_TEXTURE_MAX + 1);
+    TVector<TVector<const TVariable *>> groupedImageUniforms(HLSL_RWTEXTURE_MAX + 1);
     for (auto &uniformIt : referencedUniforms)
     {
         // Output regular uniforms. Group sampler uniforms by type.
@@ -374,7 +412,18 @@
         }
         else if (outputType == SH_HLSL_4_1_OUTPUT && IsImage(type.getBasicType()))
         {
-            imageUniformsHLSL41Output.push_back(&variable);
+            if (type.getMemoryQualifier().readonly)
+            {
+                HLSLTextureGroup group = TextureGroup(
+                    type.getBasicType(), type.getLayoutQualifier().imageInternalFormat);
+                groupedReadonlyImageUniforms[group].push_back(&variable);
+            }
+            else
+            {
+                HLSLRWTextureGroup group = RWTextureGroup(
+                    type.getBasicType(), type.getLayoutQualifier().imageInternalFormat);
+                groupedImageUniforms[group].push_back(&variable);
+            }
         }
         else
         {
@@ -421,6 +470,8 @@
     if (outputType == SH_HLSL_4_1_OUTPUT)
     {
         unsigned int groupTextureRegisterIndex = 0;
+        unsigned int groupRWTextureRegisterIndex = 0;
+        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);
         for (int groupId = HLSL_TEXTURE_MIN; groupId < HLSL_TEXTURE_MAX; ++groupId)
@@ -431,18 +482,18 @@
         }
         mSamplerCount = groupTextureRegisterIndex;
 
-        for (const TVariable *image : imageUniformsHLSL41Output)
+        for (int groupId = HLSL_TEXTURE_MIN; groupId < HLSL_TEXTURE_MAX; ++groupId)
         {
-            const TType &type          = image->getType();
-            unsigned int registerIndex = assignUniformRegister(type, image->name(), nullptr);
-            if (type.getMemoryQualifier().readonly)
-            {
-                outputHLSL4_1_FL11Texture(out, type, *image, registerIndex);
-            }
-            else
-            {
-                outputHLSL4_1_FL11RWTexture(out, type, *image, registerIndex);
-            }
+            outputHLSLReadonlyImageUniformGroup(
+                out, HLSLTextureGroup(groupId), groupedReadonlyImageUniforms[groupId],
+                &groupTextureRegisterIndex, &imageUniformGroupIndex);
+        }
+
+        for (int groupId = HLSL_RWTEXTURE_MIN; groupId < HLSL_RWTEXTURE_MAX; ++groupId)
+        {
+            outputHLSLImageUniformGroup(out, HLSLRWTextureGroup(groupId),
+                                        groupedImageUniforms[groupId], &groupRWTextureRegisterIndex,
+                                        &imageUniformGroupIndex);
         }
     }
 }
diff --git a/src/compiler/translator/UniformHLSL.h b/src/compiler/translator/UniformHLSL.h
index aa02fb1..18e1067 100644
--- a/src/compiler/translator/UniformHLSL.h
+++ b/src/compiler/translator/UniformHLSL.h
@@ -22,8 +22,7 @@
 class UniformHLSL : angle::NonCopyable
 {
   public:
-    UniformHLSL(sh::GLenum shaderType,
-                StructureHLSL *structureHLSL,
+    UniformHLSL(StructureHLSL *structureHLSL,
                 ShShaderOutput outputType,
                 const std::vector<Uniform> &uniforms,
                 unsigned int firstUniformRegister);
@@ -67,14 +66,6 @@
                                     const TType &type,
                                     const TVariable &variable,
                                     const unsigned int registerIndex);
-    void outputHLSL4_1_FL11Texture(TInfoSinkBase &out,
-                                   const TType &type,
-                                   const TVariable &variable,
-                                   const unsigned int registerIndex);
-    void outputHLSL4_1_FL11RWTexture(TInfoSinkBase &out,
-                                     const TType &type,
-                                     const TVariable &variable,
-                                     const unsigned int registerIndex);
     void outputUniform(TInfoSinkBase &out,
                        const TType &type,
                        const TVariable &variable,
@@ -95,12 +86,26 @@
         const TMap<const TVariable *, TString> &samplerInStructSymbolsToAPINames,
         unsigned int *groupTextureRegisterIndex);
 
+    void outputHLSLImageUniformIndices(TInfoSinkBase &out,
+                                       const TVector<const TVariable *> &group,
+                                       unsigned int imageArrayIndex,
+                                       unsigned int *groupRegisterCount);
+    void outputHLSLReadonlyImageUniformGroup(TInfoSinkBase &out,
+                                             const HLSLTextureGroup textureGroup,
+                                             const TVector<const TVariable *> &group,
+                                             unsigned int *groupTextureRegisterIndex,
+                                             unsigned int *imageUniformGroupIndex);
+    void outputHLSLImageUniformGroup(TInfoSinkBase &out,
+                                     const HLSLRWTextureGroup textureGroup,
+                                     const TVector<const TVariable *> &group,
+                                     unsigned int *groupTextureRegisterIndex,
+                                     unsigned int *imageUniformGroupIndex);
+
     unsigned int mUniformRegister;
     unsigned int mUniformBlockRegister;
     unsigned int mTextureRegister;
     unsigned int mRWTextureRegister;
     unsigned int mSamplerCount;
-    sh::GLenum mShaderType;
     StructureHLSL *mStructureHLSL;
     ShShaderOutput mOutputType;