WGSL: support small stride arrays in uniforms

WGSL requires arrays in the uniform address space to have a stride
a multiple of 16.

This CL makes WGSL translator emit wrapper structs for array element
types used in the uniform address space, when the array stride is
not a multiple of 16. The exception is for structs that aren't
an aligned size of 16n, and for any types matCx2, since they are
(or will be) handled in different ways that ensure alignment to 16.
This should leave only f32, i32, u32, and vec2.
See https://www.w3.org/TR/WGSL/#example-67da5de6 for an example
of using a wrapper struct.

This requires converting arrays with a wrapper struct element type
to arrays with an unwrapped element type when they are first used;
this can be "optimized" later for the common case of accessing a
single array element, which can then be unwrapped immediately. This
CL generates WGSL conversion functions when necessary.

After this, the only types that can't yet be used in a uniform
are matCx2 and bools.

This is #2 in
https://docs.google.com/document/d/17Qku1QEbLDhvJS-JJ9lPQAbnuZtLxWhG-ha5eCUhtEY/edit?tab=t.0#bookmark=id.rt3slgehd4te

Bug: angleproject:376553328
Change-Id: I1edfa7f481a6cbf5b595643aae8728e67bc4b770
Reviewed-on: https://chromium-review.googlesource.com/c/angle/angle/+/6092038
Reviewed-by: Shahbaz Youssefi <[email protected]>
Reviewed-by: Liza Burakova <[email protected]>
Reviewed-by: Matt Denton <[email protected]>
Commit-Queue: Matt Denton <[email protected]>

src/compiler/translator/Common.h

@@ -142,6 +142,20 @@ class TMap : public std::map<K, D, CMP, pool_allocator<std::pair<const K, D>>>
     {}
 };
 
+template <class K, class CMP = std::less<K>>
+class TSet : public std::set<K, CMP, pool_allocator<K>>
+{
+  public:
+    POOL_ALLOCATOR_NEW_DELETE
+    typedef pool_allocator<K> tAllocator;
+
+    TSet() : std::set<K, CMP, tAllocator>() {}
+    // use correct two-stage name lookup supported in gcc 3.4 and above
+    TSet(const tAllocator &a)
+        : std::set<K, CMP, tAllocator>(std::map<K, CMP, tAllocator>::key_compare(), a)
+    {}
+};
+
 // Basic implementation of C++20's span for use with pool-allocated containers (TVector) or static
 // arrays.  This is used by the array sizes member of TType to allow arrayed types to be
 // constexpr-constructed.

src/compiler/translator/IntermNode.cpp

@@ -20,6 +20,7 @@
 #include "common/utilities.h"
 #include "compiler/translator/Diagnostics.h"
 #include "compiler/translator/ImmutableString.h"
+#include "compiler/translator/ImmutableStringBuilder.h"
 #include "compiler/translator/IntermNode.h"
 #include "compiler/translator/SymbolTable.h"
 #include "compiler/translator/util.h"
@@ -1761,28 +1762,35 @@ bool TIntermSwizzle::offsetsMatch(int offset) const
     return mSwizzleOffsets.size() == 1 && mSwizzleOffsets[0] == offset;
 }
 
-void TIntermSwizzle::writeOffsetsAsXYZW(TInfoSinkBase *out) const
+ImmutableString TIntermSwizzle::getOffsetsAsXYZW() const
 {
+    ImmutableStringBuilder offsets(mSwizzleOffsets.size());
     for (const int offset : mSwizzleOffsets)
     {
         switch (offset)
         {
             case 0:
-                *out << "x";
+                offsets << "x";
                 break;
             case 1:
-                *out << "y";
+                offsets << "y";
                 break;
             case 2:
-                *out << "z";
+                offsets << "z";
                 break;
             case 3:
-                *out << "w";
+                offsets << "w";
                 break;
             default:
                 UNREACHABLE();
         }
     }
+    return offsets;
+}
+
+void TIntermSwizzle::writeOffsetsAsXYZW(TInfoSinkBase *out) const
+{
+    *out << getOffsetsAsXYZW();
 }
 
 TQualifier TIntermBinary::GetCommaQualifier(int shaderVersion,

src/compiler/translator/IntermNode.h

@@ -463,6 +463,7 @@ class TIntermSwizzle : public TIntermExpression
     bool hasSideEffects() const override { return mOperand->hasSideEffects(); }
 
     TIntermTyped *getOperand() { return mOperand; }
+    ImmutableString getOffsetsAsXYZW() const;
     void writeOffsetsAsXYZW(TInfoSinkBase *out) const;
 
     const TVector<int> &getSwizzleOffsets() { return mSwizzleOffsets; }

src/compiler/translator/wgsl/OutputUniformBlocks.cpp

@@ -11,6 +11,7 @@
 #include "common/utilities.h"
 #include "compiler/translator/BaseTypes.h"
 #include "compiler/translator/Compiler.h"
+#include "compiler/translator/ImmutableStringBuilder.h"
 #include "compiler/translator/InfoSink.h"
 #include "compiler/translator/IntermNode.h"
 #include "compiler/translator/SymbolUniqueId.h"
@@ -87,6 +88,59 @@ bool RecordUniformBlockMetadata(TIntermBlock *root, UniformBlockMetadata &outMet
     return true;
 }
 
+bool OutputUniformWrapperStructsAndConversions(
+    TInfoSinkBase &output,
+    const WGSLGenerationMetadataForUniforms &wgslGenerationMetadataForUniforms)
+{
+    for (const TType &type : wgslGenerationMetadataForUniforms.arrayElementTypesInUniforms)
+    {
+        // Structs don't need wrapper structs.
+        ASSERT(type.getStruct() == nullptr);
+        // Multidimensional arrays not currently supported in uniforms
+        ASSERT(!type.isArray());
+
+        output << "struct " << MakeUniformWrapperStructName(&type) << "\n{\n";
+        output << "  @align(16) " << kWrappedStructFieldName << " : ";
+        WriteWgslType(output, type, {});
+        output << "\n};\n";
+    }
+
+    for (const TType &type :
+         wgslGenerationMetadataForUniforms.arrayElementTypesThatNeedUnwrappingConversions)
+    {
+        // Should be a subset of the types that have had wrapper structs generated above, otherwise
+        // it's impossible to unwrap them!
+        TType innerType = type;
+        innerType.toArrayElementType();
+        ASSERT(wgslGenerationMetadataForUniforms.arrayElementTypesInUniforms.count(innerType) != 0);
+
+        // This could take ptr<uniform, typeName>, with the unrestricted_pointer_parameters
+        // extension. This is probably fine.
+        output << "fn " << MakeUnwrappingArrayConversionFunctionName(&type) << "(wrappedArr : ";
+        WriteWgslType(output, type, {WgslAddressSpace::Uniform});
+        output << ") -> ";
+        WriteWgslType(output, type, {WgslAddressSpace::NonUniform});
+        output << "\n{\n";
+        output << "  var retVal : ";
+        WriteWgslType(output, type, {WgslAddressSpace::NonUniform});
+        output << ";\n";
+        output << "  for (var i : u32 = 0; i < " << type.getOutermostArraySize() << "; i++) {;\n";
+        output << "    retVal[i] = wrappedArr[i]." << kWrappedStructFieldName << ";\n";
+        output << "  }\n";
+        output << "  return retVal;\n";
+        output << "}\n";
+    }
+
+    return true;
+}
+
+ImmutableString MakeUnwrappingArrayConversionFunctionName(const TType *type)
+{
+    return BuildConcatenatedImmutableString("ANGLE_Convert_", MakeUniformWrapperStructName(type),
+                                            "_ElementsTo_", type->getBuiltInTypeNameString(),
+                                            "_Elements");
+}
+
 bool OutputUniformBlocks(TCompiler *compiler, TIntermBlock *root)
 {
     // TODO(anglebug.com/42267100): This should eventually just be handled the same way as a regular
@@ -130,7 +184,7 @@ bool OutputUniformBlocks(TCompiler *compiler, TIntermBlock *root)
 
         TIntermDeclaration *declNode = globalVars.find(shaderVar.name)->second;
         const TVariable *astVar      = &ViewDeclaration(*declNode).symbol.variable();
-        WriteWgslType(output, astVar->getType());
+        WriteWgslType(output, astVar->getType(), {WgslAddressSpace::Uniform});
 
         output << ",\n";
     }

src/compiler/translator/wgsl/OutputUniformBlocks.h

@@ -7,18 +7,21 @@
 #ifndef COMPILER_TRANSLATOR_WGSL_OUTPUT_UNIFORM_BLOCKS_H_
 #define COMPILER_TRANSLATOR_WGSL_OUTPUT_UNIFORM_BLOCKS_H_
 
+#include "compiler/translator/Common.h"
 #include "compiler/translator/Compiler.h"
 #include "compiler/translator/IntermNode.h"
 
 namespace sh
 {
 
-const char kDefaultUniformBlockVarType[]     = "ANGLE_DefaultUniformBlock";
-const char kDefaultUniformBlockVarName[]     = "ANGLE_defaultUniformBlock";
-const uint32_t kDefaultUniformBlockBindGroup = 0;
+const char kDefaultUniformBlockVarType[]           = "ANGLE_DefaultUniformBlock";
+const char kDefaultUniformBlockVarName[]           = "ANGLE_defaultUniformBlock";
+const uint32_t kDefaultUniformBlockBindGroup       = 0;
 const uint32_t kDefaultVertexUniformBlockBinding   = 0;
 const uint32_t kDefaultFragmentUniformBlockBinding = 1;
 
+const char kWrappedStructFieldName[] = "elem";
+
 struct UniformBlockMetadata
 {
     // A list of structs used anywhere in the uniform address space. These will require special
@@ -33,8 +36,34 @@ struct UniformBlockMetadata
 // `outMetadata`.
 bool RecordUniformBlockMetadata(TIntermBlock *root, UniformBlockMetadata &outMetadata);
 
+// Based on the GLSL, some extra WGSL will have to be generated so it can be referenced by
+// the WGSL generated by the traverser, This tracks exactly which WGSL snippets will need to be
+// generated,
+struct WGSLGenerationMetadataForUniforms
+{
+    // Arrays must have a stride of at least 16 if used in the uniform address spaces. If the array
+    // element type doesn't have an aligned size of a multiple of 16 (e.g. f32), the element type
+    // must be wrapped in a struct which is then aligned to 16. Adding to
+    // `arrayElementTypesInUniforms` will cause `OutputUniformWrapperStructsAndConversions` to
+    // generate a WGSL wrapper struct of the form:
+    //
+    // struct ANGLE_wrapper_f32 {
+    //   @align(16) elem : f32;
+    // };
+    TSet<TType> arrayElementTypesInUniforms;
+
+    // If we need to convert arrays with wrapped element types into arrays with unwrapped element
+    // types, the necessary conversions are listed here.
+    TSet<TType> arrayElementTypesThatNeedUnwrappingConversions;
+};
+bool OutputUniformWrapperStructsAndConversions(
+    TInfoSinkBase &output,
+    const WGSLGenerationMetadataForUniforms &arrayElementTypesInUniforms);
+
+ImmutableString MakeUnwrappingArrayConversionFunctionName(const TType *type);
+
 // TODO(anglebug.com/42267100): for now does not output all uniform blocks,
-// just the default block. (fails for  matCx2, bool, and arrays with stride less than 16.)
+// just the default block. (fails for  matCx2, bool.)
 bool OutputUniformBlocks(TCompiler *compiler, TIntermBlock *root);
 
 }  // namespace sh

src/compiler/translator/wgsl/RewritePipelineVariables.cpp

@@ -352,7 +352,7 @@ class RewritePipelineVarOutputBuilder
 
             // E.g. `_uuserVar : i32,`.
             TStringStream typeStream;
-            WriteWgslType(typeStream, astVar->getType());
+            WriteWgslType(typeStream, astVar->getType(), {});
             TString type = typeStream.str();
             ImmutableString globalStructVar =
                 BuildConcatenatedImmutableString(userVarName, " : ", type.c_str(), ",");
@@ -411,11 +411,11 @@ bool RewritePipelineVarOutputBuilder::GenerateMainFunctionAndIOStructs(
 RewritePipelineVarOutput::RewritePipelineVarOutput(sh::GLenum shaderType) : mShaderType(shaderType)
 {}
 
-bool RewritePipelineVarOutput::IsInputVar(TSymbolUniqueId angleInputVar)
+bool RewritePipelineVarOutput::IsInputVar(TSymbolUniqueId angleInputVar) const
 {
     return mAngleInputVars.count(angleInputVar.get()) > 0;
 }
-bool RewritePipelineVarOutput::IsOutputVar(TSymbolUniqueId angleOutputVar)
+bool RewritePipelineVarOutput::IsOutputVar(TSymbolUniqueId angleOutputVar) const
 {
     return mAngleOutputVars.count(angleOutputVar.get()) > 0;
 }

src/compiler/translator/wgsl/RewritePipelineVariables.h

@@ -99,8 +99,8 @@ struct RewritePipelineVarOutput
 
     // Every time the translator goes to output a TVariable/TSymbol it checks these functions to see
     // if it should generate a struct access instead.
-    bool IsInputVar(TSymbolUniqueId angleInputVar);
-    bool IsOutputVar(TSymbolUniqueId angleOutputVar);
+    bool IsInputVar(TSymbolUniqueId angleInputVar) const;
+    bool IsOutputVar(TSymbolUniqueId angleOutputVar) const;
 
     bool OutputStructs(TInfoSinkBase &output);
     bool OutputMainFunction(TInfoSinkBase &output);