Spill DX12 unbounded array space assignments

DX12 binding semantics assign a register to each binding in an array and do not permit overlap. Spaces act as a "namespace" to permit multiple unbounded arrays to be present in a shader. Instead of restricting unbounded array use within an SRG, this change modifies the space assignment strategy. Unbounded arrays (on platforms that require it, DX12 in this case) are designated a unique space, starting from 1000 and counting up. This "spill space" is orthogonal to the space assignment corresponding with the SRG frequency. Signed-off-by: Jeremy Ong <jcong@amazon.com>
o3de · Mar 27, 2022 · c703adb · c703adb
1 parent ec7e239
commit c703adb
Show file tree

Hide file tree

Showing 14 changed files with 115 additions and 32 deletions.
diff --git a/.gitignore b/.gitignore
@@ -6,6 +6,7 @@ bin/*
 src/generated/*.interp
 src/generated/*.tokens
 src/generated/java/
+src/dist
 .vs
 build
 /src/.antlr/azslLexer.interp
@@ -14,4 +15,4 @@ build
 /src/.antlr/azslParser.interp
 /src/.antlr/azslParser.java
 /src/.antlr/azslParser.tokens
-/tests/testfuncs.pyc
+/tests/testfuncs.pyc
diff --git a/Platform/Windows/src/DirectX12PlatformEmitter.cpp b/Platform/Windows/src/DirectX12PlatformEmitter.cpp
@@ -71,11 +71,14 @@ namespace AZ::ShaderCompiler
             {
                 if (param.m_type != RootParamType::SrgConstantCB && param.m_type != RootParamType::RootConstantCB) // already treated in the previous loop
                 {
+                    // In DX12, unbounded arrays are assigned their own space instead of using the space assigned to the SRG
+                    int space = param.m_isUnboundedArray ? param.m_spillSpace : param.m_registerBinding.m_pair[querySet].m_logicalSpace;
+
                     std::stringstream rootParam;
                     rootParam << RootParamType::ToStr(param.m_type)
                         << "(" << ToLower(BindingType::ToStr(RootParamTypeToBindingType(param.m_type)));  // eg "CBV(b" or "SRV(t" or "Sampler(s"
                     rootParam << std::to_string(param.m_registerBinding.m_pair[querySet].m_registerIndex)
-                        << ", space = " << std::to_string(param.m_registerBinding.m_pair[querySet].m_logicalSpace)
+                        << ", space = " << std::to_string(space)
                         << ", numDescriptors = " << std::to_string(param.m_registerRange) << ")";
 
                     const auto* memberInfo = codeEmitter.GetIR()->GetSymbolSubAs<VarInfo>(param.m_uid.m_name);
@@ -125,5 +128,4 @@ namespace AZ::ShaderCompiler
 
         return Decorate("#define sig ", Join(rootAttrList.begin(), rootAttrList.end(), ", \" \\\n"), "\"\n\n");
     }
-
 }
diff --git a/Platform/Windows/src/DirectX12PlatformEmitter.h b/Platform/Windows/src/DirectX12PlatformEmitter.h
@@ -22,6 +22,8 @@ namespace AZ::ShaderCompiler
         [[nodiscard]]
         string GetRootSig(const CodeEmitter& codeEmitter, const RootSigDesc& rootSig, const Options& options, BindingPair::Set signatureQuery) const override final;
 
+        bool UnboundedArraysUseSpillSpace() const override { return true; }
+
     private:
         DirectX12PlatformEmitter() : PlatformEmitter {} {};
     };

diff --git a/src/AzslcBackend.cpp b/src/AzslcBackend.cpp
@@ -228,10 +228,17 @@ namespace AZ::ShaderCompiler
         }
     }
 
-    void MultiBindingLocationMaker::SignalRegisterIncrement(BindingType regType, int count /* = 1*/)
+    void MultiBindingLocationMaker::SignalRegisterIncrement(BindingType regType, int count, bool isUnbounded)
     {
         m_untainted.m_registerPos[regType] += count;
         m_merged.m_registerPos[regType] += count;
+
+        if (isUnbounded)
+        {
+            m_untainted.m_unboundedSpillSpace = m_unboundedSpillSpace;
+            m_merged.m_unboundedSpillSpace = m_unboundedSpillSpace;
+            ++m_unboundedSpillSpace;
+        }
     }
 
     BindingPair MultiBindingLocationMaker::GetCurrent(BindingType regType)
@@ -573,9 +580,19 @@ namespace AZ::ShaderCompiler
         }
 
         auto regType = RootParamTypeToBindingType(paramType);
-        auto srgElementDesc = RootSigDesc::SrgParamDesc{ id, paramType, bindInfo.GetCurrent(regType), count, -1, isUnboundedArray};
+
+        auto srgElementDesc = RootSigDesc::SrgParamDesc{
+            id,
+            paramType,
+            bindInfo.GetCurrent(regType),
+            count,
+            -1,
+            bindInfo.m_unboundedSpillSpace,
+            isUnboundedArray };
+
         rootSig.m_descriptorMap.emplace(id, srgElementDesc);
-        bindInfo.SignalRegisterIncrement(regType, count);
+        bindInfo.SignalRegisterIncrement(regType, count, isUnboundedArray);
+
         return srgElementDesc;
     }
 
@@ -588,7 +605,10 @@ namespace AZ::ShaderCompiler
 
     RootSigDesc Backend::BuildSignatureDescription(const Options& options, int num32BitConst) const
     {
-        MultiBindingLocationMaker bindInfo{ options };
+        // We start at an arbitrarily high number to avoid conflicts with spaces occupied by any user declarations
+        m_unboundedSpillSpace = 1000;
+
+        MultiBindingLocationMaker bindInfo{ options, m_unboundedSpillSpace };
         RootSigDesc rootSig;
 
         auto allSrgs = m_ir->m_symbols.GetOrderedSymbolsOfSubType_2<SRGInfo>();

diff --git a/src/AzslcBackend.h b/src/AzslcBackend.h
@@ -64,6 +64,9 @@ namespace AZ::ShaderCompiler
             BindingPair m_registerBinding;
             int m_registerRange = -1;
             int m_num32BitConstants = -1;
+
+            int m_spillSpace = -1;
+
             // This flag is added so m_registerRange can take the value
             // of 1 and at the same time We do not forget that m_uid refers
             // to an unbounded array.
@@ -98,6 +101,10 @@ namespace AZ::ShaderCompiler
         int GetAccumulated(BindingType r) const;
 
         int m_space = 0;  //<! logical space
+
+        // See: MultiBindingLocationMarker::m_unboundedSpillSpace
+        int m_unboundedSpillSpace;
+
         int m_registerPos[BindingType::EndEnumeratorSentinel_] = {0};   //!< register index, one by type.
         int m_accumulated[BindingType::EndEnumeratorSentinel_] = {0};  //!< Counter for total usage independently from space increments
         int m_accumulatedUnused[BindingType::EndEnumeratorSentinel_] = {0};  //!< Counter for holes created by indices unification
@@ -109,21 +116,28 @@ namespace AZ::ShaderCompiler
     class MultiBindingLocationMaker
     {
     public:
-        MultiBindingLocationMaker(const Options& options)
-            : m_options(options)
+        MultiBindingLocationMaker(const Options& options, int& unboundedSpillSpace)
+            : m_options{ options }
+            , m_unboundedSpillSpace{ unboundedSpillSpace }
         {}
 
         void SignalIncrementSpace(std::function<void(int, int)> warningMessageFunctionForMinDescOvershoot);
 
         void SignalUnifyIndices();
 
-        void SignalRegisterIncrement(BindingType regType, int count = 1);
+        void SignalRegisterIncrement(BindingType regType, int count, bool isUnbounded);
 
         BindingPair GetCurrent(BindingType regType);
 
         SingleBindingLocationTracker m_untainted;
         SingleBindingLocationTracker m_merged;
         Options m_options;
+
+        // On some platforms (DX12), descriptor arrays occupy an individual register slot, and spaces are used
+        // to prevent overlapping ranges. When an unbounded array is encountered, we immediately assign it to
+        // the value of this member variable and increment. This is initialized in the constructor because the
+        // space we spill to must not collide with any other SRG declared in the shader.
+        int& m_unboundedSpillSpace;
     };
 
     //! This class intends to be a base umbrella for compiler back-end services.
@@ -174,6 +188,9 @@ namespace AZ::ShaderCompiler
         std::ostream&               m_out;
         IntermediateRepresentation* m_ir;
         TokenStream*                m_tokens;
+
+        // See MultiBindingLocationMaker::m_unboundedSpillSpace
+        mutable int m_unboundedSpillSpace;
     };
 
     // independent utility functions

diff --git a/src/AzslcEmitter.cpp b/src/AzslcEmitter.cpp
@@ -1043,7 +1043,7 @@ namespace AZ::ShaderCompiler
         // note: instead of redoing this work ad-hoc, EmitText could be used directly on the ext type.
         const auto genericType = "<" + GetTranslatedName(varInfo->m_typeInfoExt.m_genericParameter, UsageContext::ReferenceSite) + ">";
 
-        const string spaceX = (options.m_useLogicalSpaces) ? ", space" + std::to_string(bindInfo.m_registerBinding.m_pair[bindSet].m_logicalSpace) : "";
+        const string spaceX = (options.m_useLogicalSpaces) ? ", space" + std::to_string(ResolveBindingSpace(bindInfo, bindSet)) : "";
         m_out << "ConstantBuffer " << genericType << " " << cbName;
         if (bindInfo.m_isUnboundedArray)
         {
@@ -1066,7 +1066,7 @@ namespace AZ::ShaderCompiler
 
         EmitEmptyLinesToLineNumber(varInfo->GetOriginalLineNumber());
 
-        const string spaceX = (options.m_useLogicalSpaces) ? ", space" + std::to_string(bindInfo.m_registerBinding.m_pair[bindSet].m_logicalSpace) : "";
+        const string spaceX = (options.m_useLogicalSpaces) ? ", space" + std::to_string(ResolveBindingSpace(bindInfo, bindSet)) : "";
         m_out << (varInfo->m_samplerState->m_isComparison ? "SamplerComparisonState " : "SamplerState ")
               << ReplaceSeparators(sId.m_name, Underscore);
         if (bindInfo.m_isUnboundedArray)
@@ -1112,9 +1112,10 @@ namespace AZ::ShaderCompiler
         string varType = GetTranslatedName(varInfo->m_typeInfoExt, UsageContext::DeclarationSite);
         auto   registerTypeLetter = ToLower(BindingType::ToStr(RootParamTypeToBindingType(bindInfo.m_type)));
         optional<string> stringifiedLogicalSpace;
+
         if (options.m_useLogicalSpaces)
         {
-            stringifiedLogicalSpace = std::to_string(bindInfo.m_registerBinding.m_pair[bindSet].m_logicalSpace);
+            stringifiedLogicalSpace = std::to_string(ResolveBindingSpace(bindInfo, bindSet));
         }
 
         EmitEmptyLinesToLineNumber(varInfo->GetOriginalLineNumber());
@@ -1363,4 +1364,17 @@ namespace AZ::ShaderCompiler
             m_out << "\n";
         }
     }
+
+    int CodeEmitter::ResolveBindingSpace(const RootSigDesc::SrgParamDesc& bindInfo, BindingPair::Set bindSet) const
+    {
+        if (bindInfo.m_isUnboundedArray && GetPlatformEmitter().UnboundedArraysUseSpillSpace())
+        {
+            return bindInfo.m_spillSpace;
+        }
+        else
+        {
+            return bindInfo.m_registerBinding.m_pair[bindSet].m_logicalSpace;
+        }
+    }
+
 }
diff --git a/src/AzslcEmitter.h b/src/AzslcEmitter.h
@@ -228,5 +228,8 @@ namespace AZ::ShaderCompiler
         // whether We are using a regular std::ostream or an instance of NewLineCounterStream.
         template <class StreamLike>
         void GetTextInStreamInternal(misc::Interval interval, StreamLike& output, bool emitNewLines) const;
+
+        // Given an SRG parameter, determines the space it belongs to based on the platform
+        int ResolveBindingSpace(const RootSigDesc::SrgParamDesc& bindInfo, BindingPair::Set bindSet) const;
     };
 }
diff --git a/src/AzslcMain.cpp b/src/AzslcMain.cpp
@@ -526,13 +526,6 @@ int main(int argc, const char* argv[])
             std::for_each(namespaces.begin(), namespaces.end(),
                 [&](const string& space) { ir.AddAttributeNamespaceFilter(space); });
 
-            UnboundedArraysValidator::Options unboundedArraysValidationOptions = { useSpaces, uniqueIdx };
-            if (*maxSpacesOpt)
-            {
-                unboundedArraysValidationOptions.m_maxSpaces = maxSpaces;
-            }
-            ir.m_sema.m_unboundedArraysValidator.SetOptions(unboundedArraysValidationOptions);
-
             // semantic logic and validation
             walker.walk(&semanticListener, tree);
 

diff --git a/src/AzslcPlatformEmitter.h b/src/AzslcPlatformEmitter.h
@@ -74,5 +74,7 @@ namespace AZ::ShaderCompiler
         //! Aligns the size for the data containing the root constants.
         //! @param size  The size of stride
         virtual uint32_t AlignRootConstants(uint32_t size) const;
+
+        virtual bool UnboundedArraysUseSpillSpace() const { return false; }
     };
 }
diff --git a/src/AzslcReflection.cpp b/src/AzslcReflection.cpp
@@ -631,13 +631,23 @@ namespace AZ::ShaderCompiler
         m_out << GetVariantList(options);
     }
 
-    static void ReflectBinding(Json::Value& output, const RootSigDesc::SrgParamDesc& bindInfo)
+    void CodeReflection::ReflectBinding(Json::Value& output, const RootSigDesc::SrgParamDesc& bindInfo) const
     {
         output["count"] = (bindInfo.m_isUnboundedArray) ? -1 : bindInfo.m_registerRange;
         output["index"] = bindInfo.m_registerBinding.m_pair[BindingPair::Set::Untainted].m_registerIndex;
-        output["space"] = bindInfo.m_registerBinding.m_pair[BindingPair::Set::Untainted].m_logicalSpace;
         output["index-merged"] = bindInfo.m_registerBinding.m_pair[BindingPair::Set::Merged].m_registerIndex;
-        output["space-merged"] = bindInfo.m_registerBinding.m_pair[BindingPair::Set::Merged].m_logicalSpace;
+
+        if (bindInfo.m_isUnboundedArray && GetPlatformEmitter().UnboundedArraysUseSpillSpace())
+        {
+            output["space"] = bindInfo.m_spillSpace;
+            output["space-merged"] = bindInfo.m_spillSpace;
+        }
+        else
+        {
+            output["space"] = bindInfo.m_registerBinding.m_pair[BindingPair::Set::Untainted].m_logicalSpace;
+            output["space-merged"] = bindInfo.m_registerBinding.m_pair[BindingPair::Set::Merged].m_logicalSpace;
+        }
+
     }
 
     void CodeReflection::DumpSRGLayout(const Options& options) const
@@ -724,6 +734,7 @@ namespace AZ::ShaderCompiler
                 dataView["id"]     = ExtractLeaf(tId.m_name).data();
                 dataView["type"]   = viewName;
                 dataView["usage"]  = (isReadWriteView) ? "ReadWrite" : "Read";
+
                 ReflectBinding(dataView, bindInfo);
                 dataView["stride"] = strideSize;
 
@@ -946,7 +957,7 @@ namespace AZ::ShaderCompiler
                     return allEntriesJson;
                 };
 
-                auto makeJsonNodeForOneResource = [&dependencyListToJson](const set<IdentifierUID>& dependencyList,
+                auto makeJsonNodeForOneResource = [this, &dependencyListToJson](const set<IdentifierUID>& dependencyList,
                     const RootSigDesc::SrgParamDesc& binding,
                     const Json::Value& allConstants)
                 {

diff --git a/src/AzslcReflection.h b/src/AzslcReflection.h
@@ -43,6 +43,8 @@ namespace AZ::ShaderCompiler
         void DumpResourceBindingDependencies(const Options& options) const;
 
     private:
+        // Serialize a binding to Json
+        void ReflectBinding(Json::Value& output, const RootSigDesc::SrgParamDesc& bindInfo) const;
 
         //! Builds member variable packing information and adds it to the membersContainer
         uint32_t BuildMemberLayout(Json::Value& membersContainer, string_view namePrefix, const IdentifierUID& memberId, const bool isArrayItr, const Options& options, const AZ::ShaderCompiler::Packing::Layout layoutPacking, uint32_t& offset) const;

diff --git a/src/AzslcSemanticOrchestrator.cpp b/src/AzslcSemanticOrchestrator.cpp
@@ -690,12 +690,6 @@ namespace AZ::ShaderCompiler
         TypeClass typeClass = varInfo.GetTypeClass();
         assert(typeClass != TypeClass::Alias);
 
-        string errorMessage;
-        if (!m_unboundedArraysValidator.CheckFieldCanBeAddedToSrg(isUnboundedArray, srgUid, srgInfo, varUid, varInfo, typeClass, &errorMessage))
-        {
-            ThrowAzslcOrchestratorException(ORCHESTRATOR_UNBOUNDED_RESOURCE_ISSUE, ctx->start, errorMessage);
-        }
-
         if (typeClass == TypeClass::ConstantBuffer)
         {
             srgInfo.m_CBs.push_back(varUid);

diff --git a/src/AzslcSemanticOrchestrator.h b/src/AzslcSemanticOrchestrator.h
@@ -9,7 +9,6 @@
 
 #include "AzslcScopeTracker.h"
 #include "PreprocessorLineDirectiveFinder.h"
-#include "AzslcUnboundedArraysValidator.h"
 
 namespace AZ::ShaderCompiler
 {
@@ -423,7 +422,6 @@ namespace AZ::ShaderCompiler
         ScopeTracker*     m_scope;
         azslLexer*        m_lexer;
         PreprocessorLineDirectiveFinder* m_preprocessorLineDirectiveFinder;
-        UnboundedArraysValidator m_unboundedArraysValidator;
 
         //! cached property informing of the presence of at least one input attachment use.
         bool              m_subpassInputSeen = false;

diff --git a/tests/Samples/UnboundedArrays2.azsl_manual b/tests/Samples/UnboundedArrays2.azsl_manual
@@ -0,0 +1,24 @@
+ShaderResourceGroupSemantic slot1
+{
+    FrequencyId = 1;
+};
+
+struct MyStruct
+{
+    float4 m_a;
+    float4 m_b;
+};
+
+
+ShaderResourceGroup SRG1 : slot1
+{
+
+    Texture2D<float4>        m_texSRVa[];      // t0+, space1000
+    Texture2D<float4>        m_texSRVb[];      // t0+, space1001
+    RWTexture2D<float4>      m_texUAVa[];      // u0+, space1002
+    RWTexture2D<float4>      m_texUAVb[];      // u0+, space1003
+    Sampler                  m_samplera[];     // s0+, space1004
+    Sampler                  m_samplerb[];     // s0+, space1005
+    ConstantBuffer<MyStruct> m_structArraya[]; // b0+, space1006
+    ConstantBuffer<MyStruct> m_structArrayb[]; // b0+, space1007
+};