Merge pull request #6905 from aws-lumberyard-dev/Atom/dmcdiar/ATOM-15813

Allow InstanceDatabase<Shader> to support multiple instances with different supervariants

Gems/Atom/RPI/Code/Include/Atom/RPI.Public/Shader/Shader.h

@@ -136,11 +136,6 @@ namespace AZ
             //! This tag corresponds to the ShaderAsset object's DrawListName.
             RHI::DrawListTag GetDrawListTag() const;
 
-            //! Changes the supervariant of the shader to the specified supervariantIndex.
-            //! [GFX TODO][ATOM-15813]: this can be removed when the shader InstanceDatabase can support multiple shader
-            //! instances with different supervariants.
-            void ChangeSupervariant(SupervariantIndex supervariantIndex);
-
         private:
             explicit Shader(const SupervariantIndex& supervariantIndex) : m_supervariantIndex(supervariantIndex){};
             Shader() = delete;

Gems/Atom/RPI/Code/Include/Atom/RPI.Reflect/Shader/ShaderVariantAsset.h

@@ -50,6 +50,7 @@ namespace AZ
             RPI::ShaderVariantStableId GetStableId() const { return m_stableId;  }
 
             const ShaderVariantId& GetShaderVariantId() const { return m_shaderVariantId; }
+            uint32_t GetSupervariantIndex() const;
 
             //! Returns the shader stage function associated with the provided stage enum value.
             const RHI::ShaderStageFunction* GetShaderStageFunction(RHI::ShaderStage shaderStage) const;

Gems/Atom/RPI/Code/Source/RPI.Public/Shader/Shader.cpp

@@ -26,21 +26,34 @@ namespace AZ
         Data::Instance<Shader> Shader::FindOrCreate(const Data::Asset<ShaderAsset>& shaderAsset, const Name& supervariantName)
         {
             auto anySupervariantName = AZStd::any(supervariantName);
-            Data::Instance<Shader> shaderInstance = Data::InstanceDatabase<Shader>::Instance().FindOrCreate(
-                Data::InstanceId::CreateFromAssetId(shaderAsset.GetId()), shaderAsset, &anySupervariantName);
 
-            if (shaderInstance)
+            // retrieve the supervariant index from the shader asset
+            SupervariantIndex supervariantIndex = shaderAsset->GetSupervariantIndex(supervariantName);
+            if (!supervariantIndex.IsValid())
             {
-                // [GFX TODO][ATOM-15813] Change InstanceDatabase<Shader> to support multiple instances with different supervariants.
-                // At this time we do not support multiple supervariants loaded for a shader asset simultaneously, so if this shader
-                // is referring to the wrong supervariant we need to change it to the correct one.
-                SupervariantIndex supervariantIndex = shaderAsset->GetSupervariantIndex(supervariantName);
-                if (supervariantIndex.IsValid() && shaderInstance->GetSupervariantIndex() != supervariantIndex)
-                {
-                    shaderInstance->ChangeSupervariant(supervariantIndex);
-                }
+                AZ_Error("Shader", false, "Supervariant with name %s, was not found in shader %s", supervariantName.GetCStr(), shaderAsset->GetName().GetCStr());
+                return nullptr;
             }
 
+            // create the InstanceId from the combined assetId and supervariantIndex
+            const Data::AssetId& assetId = shaderAsset.GetId();
+            uint32_t shaderSupervariantIndex = supervariantIndex.GetIndex();
+
+            const uint32_t instanceIdDataSize = sizeof(assetId.m_guid) + sizeof(assetId.m_subId) + sizeof(shaderSupervariantIndex);
+            uint8_t instanceIdData[instanceIdDataSize];
+            uint8_t* instanceIdDataPtr = instanceIdData;
+
+            memcpy(instanceIdDataPtr, &assetId.m_guid, sizeof(assetId.m_guid));
+            instanceIdDataPtr += sizeof(assetId.m_guid);
+            memcpy(instanceIdDataPtr, &assetId.m_subId, sizeof(assetId.m_subId));
+            instanceIdDataPtr += sizeof(assetId.m_subId);
+            memcpy(instanceIdDataPtr, &shaderSupervariantIndex, sizeof(shaderSupervariantIndex));
+
+            Data::InstanceId instanceId = Data::InstanceId::CreateData(instanceIdData, instanceIdDataSize);
+
+            // retrieve the shader instance from the Instance database
+            Data::Instance<Shader> shaderInstance = Data::InstanceDatabase<Shader>::Instance().FindOrCreate(instanceId, shaderAsset, &anySupervariantName);
+
             return shaderInstance;
         }
 
@@ -248,6 +261,12 @@ namespace AZ
             AZ_Assert(shaderVariantAsset, "Reloaded ShaderVariantAsset is null");
             const ShaderVariantStableId stableId = shaderVariantAsset->GetStableId();
 
+            // check the supervariantIndex of the ShaderVariantAsset to make sure it matches the supervariantIndex of this shader instance
+            if (shaderVariantAsset->GetSupervariantIndex() != m_supervariantIndex.GetIndex())
+            {
+                return;
+            }
+
             // We make a copy of the updated variant because OnShaderVariantReinitialized must not be called inside
             // m_variantCacheMutex or deadlocks may occur.
             // Or if there is an error, we leave this object in its default state to indicate there was an error.
@@ -486,14 +505,5 @@ namespace AZ
             return m_drawListTag;
         }
 
-        void Shader::ChangeSupervariant(SupervariantIndex supervariantIndex)
-        {
-            if (supervariantIndex != m_supervariantIndex)
-            {
-                m_supervariantIndex = supervariantIndex;
-                Init(*m_asset);
-            }
-        }
-
     } // namespace RPI
 } // namespace AZ

Gems/Atom/RPI/Code/Source/RPI.Reflect/Shader/ShaderVariantAsset.cpp

@@ -65,6 +65,11 @@ namespace AZ
             return m_buildTimestamp;
         }
 
+        uint32_t ShaderVariantAsset::GetSupervariantIndex() const
+        {
+            return (m_assetId.m_subId >> SupervariantIndexBitPosition) & SupervariantIndexMaxValue;
+        }
+
         const RHI::ShaderStageFunction* ShaderVariantAsset::GetShaderStageFunction(RHI::ShaderStage shaderStage) const
         {
             return m_functionsByStage[static_cast<size_t>(shaderStage)].get();