RayTracingFeatureProcessor.h 17 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308
  1. /*
  2. * Copyright (c) Contributors to the Open 3D Engine Project.
  3. * For complete copyright and license terms please see the LICENSE at the root of this distribution.
  4. *
  5. * SPDX-License-Identifier: Apache-2.0 OR MIT
  6. *
  7. */
  8. #pragma once
  9. #include <Atom/Feature/RayTracing/RayTracingFeatureProcessorInterface.h>
  10. #include <Atom/Feature/RayTracing/RayTracingIndexList.h>
  11. #include <Atom/Feature/TransformService/TransformServiceFeatureProcessorInterface.h>
  12. #include <Atom/RHI/DeviceBufferView.h>
  13. #include <Atom/RHI/DeviceImageView.h>
  14. #include <Atom/RHI/IndexBufferView.h>
  15. #include <Atom/RHI/RayTracingAccelerationStructure.h>
  16. #include <Atom/RHI/RayTracingBufferPools.h>
  17. #include <Atom/RHI/RayTracingCompactionQueryPool.h>
  18. #include <Atom/RHI/StreamBufferView.h>
  19. #include <Atom/RPI.Public/Buffer/RingBuffer.h>
  20. #include <Atom/RPI.Public/Shader/Shader.h>
  21. #include <Atom/RPI.Reflect/Image/Image.h>
  22. #include <Atom/Utils/StableDynamicArray.h>
  23. #include <AzCore/Math/Aabb.h>
  24. #include <AzCore/Math/Color.h>
  25. #include <AzCore/Math/Transform.h>
  26. #include <RayTracing/RayTracingResourceList.h>
  27. // this define specifies that the mesh buffers and material textures are stored in the Bindless Srg
  28. // Note1: The previous implementation using separate unbounded arrays is preserved since it demonstrates a TDR caused by
  29. // the RHI unbounded array allocation. This define and the previous codepath can be removed once the TDR is
  30. // investigated and resolved.
  31. // Note2: There are corresponding USE_BINDLESS_SRG defines in the RayTracingSceneSrg.azsli and RayTracingMaterialSrg.azsli
  32. // shader files that must match the setting of this define.
  33. #define USE_BINDLESS_SRG 1
  34. namespace AZ
  35. {
  36. namespace Render
  37. {
  38. //! This feature processor manages ray tracing data for a Scene
  39. class RayTracingFeatureProcessor
  40. : public RayTracingFeatureProcessorInterface
  41. {
  42. public:
  43. AZ_CLASS_ALLOCATOR(RayTracingFeatureProcessor, AZ::SystemAllocator)
  44. AZ_RTTI(AZ::Render::RayTracingFeatureProcessor, "{5017EFD3-A996-44B0-9ED2-C47609A2EE8D}", AZ::Render::RayTracingFeatureProcessorInterface);
  45. static void Reflect(AZ::ReflectContext* context);
  46. RayTracingFeatureProcessor() = default;
  47. virtual ~RayTracingFeatureProcessor() = default;
  48. // FeatureProcessor overrides
  49. void Activate() override;
  50. void Deactivate() override;
  51. void OnRenderPipelineChanged(RPI::RenderPipeline* renderPipeline, RPI::SceneNotification::RenderPipelineChangeType changeType) override;
  52. // RayTracingFeatureProcessorInterface overrides
  53. ProceduralGeometryTypeHandle RegisterProceduralGeometryType(
  54. const AZStd::string& name,
  55. const Data::Instance<RPI::Shader>& intersectionShader,
  56. const AZStd::string& intersectionShaderName,
  57. const AZStd::unordered_map<int, uint32_t>& bindlessBufferIndices = {}) override;
  58. void SetProceduralGeometryTypeBindlessBufferIndex(
  59. ProceduralGeometryTypeWeakHandle geometryTypeHandle, const AZStd::unordered_map<int, uint32_t>& bindlessBufferIndices) override;
  60. void AddProceduralGeometry(
  61. ProceduralGeometryTypeWeakHandle geometryTypeHandle,
  62. const Uuid& uuid,
  63. const Aabb& aabb,
  64. const SubMeshMaterial& material,
  65. RHI::RayTracingAccelerationStructureInstanceInclusionMask instanceMask,
  66. uint32_t localInstanceIndex) override;
  67. void SetProceduralGeometryTransform(
  68. const Uuid& uuid, const Transform& transform, const Vector3& nonUniformScale = Vector3::CreateOne()) override;
  69. void SetProceduralGeometryLocalInstanceIndex(const Uuid& uuid, uint32_t localInstanceIndex) override;
  70. void SetProceduralGeometryMaterial(const Uuid& uuid, const SubMeshMaterial& material) override;
  71. void RemoveProceduralGeometry(const Uuid& uuid) override;
  72. int GetProceduralGeometryCount(ProceduralGeometryTypeWeakHandle geometryTypeHandle) const override;
  73. const ProceduralGeometryTypeList& GetProceduralGeometryTypes() const override { return m_proceduralGeometryTypes; }
  74. const ProceduralGeometryList& GetProceduralGeometries() const override { return m_proceduralGeometry; }
  75. void AddMesh(const AZ::Uuid& uuid, const Mesh& rayTracingMesh, const SubMeshVector& subMeshes) override;
  76. void RemoveMesh(const AZ::Uuid& uuid) override;
  77. void SetMeshTransform(const AZ::Uuid& uuid, const AZ::Transform transform, const AZ::Vector3 nonUniformScale) override;
  78. void SetMeshReflectionProbe(const AZ::Uuid& uuid, const Mesh::ReflectionProbe& reflectionProbe) override;
  79. void SetMeshMaterials(const AZ::Uuid& uuid, const SubMeshMaterialVector& subMeshMaterials) override;
  80. const SubMeshVector& GetSubMeshes() const override { return m_subMeshes; }
  81. SubMeshVector& GetSubMeshes() override { return m_subMeshes; }
  82. const MeshMap& GetMeshMap() override { return m_meshes; }
  83. uint32_t GetSubMeshCount() const override { return m_subMeshCount; }
  84. uint32_t GetSkinnedMeshCount() const override { return m_skinnedMeshCount; }
  85. Data::Instance<RPI::ShaderResourceGroup> GetRayTracingSceneSrg() const override { return m_rayTracingSceneSrg; }
  86. Data::Instance<RPI::ShaderResourceGroup> GetRayTracingMaterialSrg() const override { return m_rayTracingMaterialSrg; }
  87. const Data::Instance<RPI::Buffer> GetMeshInfoGpuBuffer() const override { return m_meshInfoGpuBuffer.GetCurrentBuffer(); }
  88. const Data::Instance<RPI::Buffer> GetMaterialInfoGpuBuffer() const override { return m_materialInfoGpuBuffer.GetCurrentBuffer(); }
  89. void Render(const RenderPacket&) override;
  90. void BeginFrame() override;
  91. uint32_t GetRevision() const override { return m_revision; }
  92. uint32_t GetProceduralGeometryTypeRevision() const override { return m_proceduralGeometryTypeRevision; }
  93. RHI::RayTracingBufferPools& GetBufferPools() override { return *m_bufferPools; }
  94. void UpdateRayTracingSrgs() override;
  95. const RHI::Ptr<RHI::RayTracingTlas>& GetTlas() const override { return m_tlas; }
  96. RHI::Ptr<RHI::RayTracingTlas>& GetTlas() override{ return m_tlas; }
  97. RHI::AttachmentId GetTlasAttachmentId() const override { return m_tlasAttachmentId; }
  98. BlasInstanceMap& GetBlasInstances() override { return m_blasInstanceMap; }
  99. AZStd::mutex& GetBlasBuiltMutex() override { return m_blasBuiltMutex; }
  100. BlasBuildList& GetBlasBuildList(int deviceIndex) override { return m_blasToBuild[deviceIndex]; }
  101. const BlasBuildList& GetSkinnedMeshBlasList() override { return m_skinnedBlasIds; };
  102. BlasBuildList& GetBlasCompactionList(int deviceIndex) override { return m_blasToCompact[deviceIndex]; }
  103. const void MarkBlasInstanceForCompaction(int deviceIndex, Data::AssetId assetId) override;
  104. const void MarkBlasInstanceAsCompactionEnqueued(int deviceIndex, Data::AssetId assetId) override;
  105. bool HasMeshGeometry() const override { return m_subMeshCount != 0; }
  106. bool HasProceduralGeometry() const override { return !m_proceduralGeometry.empty(); }
  107. bool HasGeometry() const override { return HasMeshGeometry() || HasProceduralGeometry(); }
  108. RHI::MultiDevice::DeviceMask GetDeviceMask() const override { return m_deviceMask; }
  109. private:
  110. AZ_DISABLE_COPY_MOVE(RayTracingFeatureProcessor);
  111. void UpdateBlasInstances();
  112. void UpdateMeshInfoBuffer();
  113. void UpdateProceduralGeometryInfoBuffer();
  114. void UpdateMaterialInfoBuffer();
  115. void UpdateIndexLists();
  116. void UpdateRayTracingSceneSrg();
  117. void UpdateRayTracingMaterialSrg();
  118. void RemoveBlasInstance(Data::AssetId id);
  119. static RHI::RayTracingAccelerationStructureBuildFlags CreateRayTracingAccelerationStructureBuildFlags(bool isSkinnedMesh);
  120. // flag indicating if RayTracing is enabled, currently based on device support
  121. bool m_rayTracingEnabled = false;
  122. // mesh data for meshes that should be included in ray tracing operations,
  123. // this is a map of the mesh UUID to the ray tracing data for the sub-meshes
  124. MeshMap m_meshes;
  125. SubMeshVector m_subMeshes;
  126. // buffer pools used in ray tracing operations
  127. RHI::Ptr<RHI::RayTracingBufferPools> m_bufferPools;
  128. // ray tracing acceleration structure (TLAS)
  129. RHI::Ptr<RHI::RayTracingTlas> m_tlas;
  130. // RayTracingScene and RayTracingMaterial asset and Srgs
  131. Data::Asset<RPI::ShaderAsset> m_rayTracingSrgAsset;
  132. Data::Instance<RPI::ShaderResourceGroup> m_rayTracingSceneSrg;
  133. Data::Instance<RPI::ShaderResourceGroup> m_rayTracingMaterialSrg;
  134. // current revision number of ray tracing data
  135. uint32_t m_revision = 0;
  136. // latest tlas revision number
  137. uint32_t m_tlasRevision = 0;
  138. uint32_t m_proceduralGeometryTypeRevision = 0;
  139. // total number of ray tracing sub-meshes
  140. uint32_t m_subMeshCount = 0;
  141. // TLAS attachmentId
  142. RHI::AttachmentId m_tlasAttachmentId;
  143. // cached TransformServiceFeatureProcessor
  144. TransformServiceFeatureProcessorInterface* m_transformServiceFeatureProcessor = nullptr;
  145. // mutex for the mesh and BLAS lists
  146. AZStd::mutex m_mutex;
  147. // mutex for the m_blasBuilt flag manipulation
  148. AZStd::mutex m_blasBuiltMutex;
  149. // structure for data in the m_meshInfoBuffer, shaders that use the buffer must match this type
  150. struct MeshInfo
  151. {
  152. // byte offsets into the mesh buffer views
  153. uint32_t m_indexByteOffset = 0;
  154. uint32_t m_positionByteOffset = 0;
  155. uint32_t m_normalByteOffset = 0;
  156. uint32_t m_tangentByteOffset = 0;
  157. uint32_t m_bitangentByteOffset = 0;
  158. uint32_t m_uvByteOffset = 0;
  159. RayTracingSubMeshBufferFlags m_bufferFlags = RayTracingSubMeshBufferFlags::None;
  160. uint32_t m_bufferStartIndex = 0;
  161. AZStd::array<float, 12> m_worldInvTranspose; // float3x4
  162. };
  163. // vector of MeshInfo, transferred to the meshInfoGpuBuffer
  164. using MeshInfoVector = AZStd::vector<MeshInfo>;
  165. AZStd::unordered_map<int, MeshInfoVector> m_meshInfos;
  166. RPI::RingBuffer m_meshInfoGpuBuffer{ "RayTracingMeshInfo", RPI::CommonBufferPoolType::ReadOnly, sizeof(MeshInfo) };
  167. RPI::RingBuffer m_proceduralGeometryInfoGpuBuffer{ "ProceduralGeometryInfo", RPI::CommonBufferPoolType::ReadOnly, RHI::Format::R32G32_UINT };
  168. // structure for data in the m_materialInfoBuffer, shaders that use the buffer must match this type
  169. struct alignas(16) MaterialInfo
  170. {
  171. AZStd::array<float, 4> m_baseColor; // float4
  172. AZStd::array<float, 4> m_irradianceColor; // float4
  173. AZStd::array<float, 3> m_emissiveColor; // float3
  174. float m_metallicFactor = 0.0f;
  175. float m_roughnessFactor = 0.0f;
  176. RayTracingSubMeshTextureFlags m_textureFlags = RayTracingSubMeshTextureFlags::None;
  177. uint32_t m_textureStartIndex = InvalidIndex;
  178. uint32_t m_reflectionProbeCubeMapIndex = InvalidIndex;
  179. // reflection probe data, must match the structure in ReflectionProbeData.azlsi
  180. struct alignas(16) ReflectionProbeData
  181. {
  182. AZStd::array<float, 12> m_modelToWorld; // float3x4
  183. AZStd::array<float, 12> m_modelToWorldInverse; // float3x4
  184. AZStd::array<float, 3> m_outerObbHalfLengths; // float3
  185. float m_exposure = 0.0f;
  186. AZStd::array<float, 3> m_innerObbHalfLengths; // float3
  187. uint32_t m_useReflectionProbe = 0;
  188. uint32_t m_useParallaxCorrection = 0;
  189. AZStd::array<float, 3> m_padding;
  190. };
  191. ReflectionProbeData m_reflectionProbeData;
  192. };
  193. // vector of MaterialInfo, transferred to the materialInfoGpuBuffer
  194. using MaterialInfoVector = AZStd::vector<MaterialInfo>;
  195. AZStd::unordered_map<int, MaterialInfoVector> m_materialInfos;
  196. AZStd::unordered_map<int, MaterialInfoVector> m_proceduralGeometryMaterialInfos;
  197. RPI::RingBuffer m_materialInfoGpuBuffer{ "RayTracingMaterialInfo", RPI::CommonBufferPoolType::ReadOnly, sizeof(MaterialInfo) };
  198. // update flags
  199. bool m_meshInfoBufferNeedsUpdate = false;
  200. bool m_proceduralGeometryInfoBufferNeedsUpdate = false;
  201. bool m_materialInfoBufferNeedsUpdate = false;
  202. bool m_indexListNeedsUpdate = false;
  203. // side list for looking up existing BLAS objects so they can be re-used when the same mesh is added multiple times
  204. BlasInstanceMap m_blasInstanceMap;
  205. BlasBuildList m_blasToCreate;
  206. AZStd::unordered_map<int, BlasBuildList> m_blasToBuild;
  207. AZStd::unordered_map<int, BlasBuildList> m_blasToCompact;
  208. BlasBuildList m_skinnedBlasIds;
  209. // Mutex for the queues that are modified by the RaytracingAccelerationStructurePasses
  210. // The BuildCommandList of these passes are on different threads so we need lock access to the queues
  211. AZStd::mutex m_queueMutex;
  212. struct BlasFrameEvent
  213. {
  214. int m_frameIndex = -1;
  215. RHI::MultiDevice::DeviceMask m_deviceMask;
  216. };
  217. // List of Blas instances that are enqueued for compaction
  218. // The frame index corresponds to the frame where the compaction query is ready
  219. AZStd::unordered_map<Data::AssetId, BlasFrameEvent> m_blasEnqueuedForCompact;
  220. // List of Blas instances where the compacted Blas is already built
  221. // The frame index corresponds to the frame where the uncompacted Blas is no longer needed and can be deleted
  222. AZStd::unordered_map<Data::AssetId, BlasFrameEvent> m_uncompactedBlasEnqueuedForDeletion;
  223. int m_frameIndex = 0;
  224. int m_updatedFrameIndex = 0;
  225. #if !USE_BINDLESS_SRG
  226. // Mesh buffer and material texture resources are managed with a RayTracingResourceList, which contains an internal
  227. // indirection list. This allows resource entries to be swapped inside the RayTracingResourceList when removing entries,
  228. // without invalidating the indices held here in the m_meshBufferIndices and m_materialTextureIndices lists.
  229. // mesh buffer and material texture resource lists, accessed by the shader through an unbounded array
  230. RayTracingResourceList<RHI::BufferView> m_meshBuffers;
  231. RayTracingResourceList<const RHI::ImageView> m_materialTextures;
  232. #endif
  233. // RayTracingIndexList implements an internal freelist chain stored inside the list itself, allowing entries to be
  234. // reused after elements are removed.
  235. // mesh buffer and material texture index lists, which contain the array indices of the mesh resources
  236. static const uint32_t NumMeshBuffersPerMesh = 6;
  237. AZStd::unordered_map<int, RayTracingIndexList<NumMeshBuffersPerMesh>> m_meshBufferIndices;
  238. static const uint32_t NumMaterialTexturesPerMesh = 5;
  239. AZStd::unordered_map<int, RayTracingIndexList<NumMaterialTexturesPerMesh>> m_materialTextureIndices;
  240. // Gpu buffers for the mesh and material index lists
  241. RPI::RingBuffer m_meshBufferIndicesGpuBuffer{ "RayTracingMeshBufferIndices", RPI::CommonBufferPoolType::ReadOnly, RHI::Format::R32_UINT };
  242. RPI::RingBuffer m_materialTextureIndicesGpuBuffer{ "RayTracingMaterialTextureIndices", RPI::CommonBufferPoolType::ReadOnly, RHI::Format::R32_UINT };
  243. uint32_t m_skinnedMeshCount = 0;
  244. ProceduralGeometryTypeList m_proceduralGeometryTypes;
  245. ProceduralGeometryList m_proceduralGeometry;
  246. AZStd::unordered_map<Uuid, size_t> m_proceduralGeometryLookup;
  247. RHI::Ptr<RHI::RayTracingCompactionQueryPool> m_compactionQueryPool;
  248. RHI::MultiDevice::DeviceMask m_deviceMask = {};
  249. void ConvertMaterial(MaterialInfo& materialInfo, const SubMeshMaterial& subMeshMaterial, int deviceIndex);
  250. };
  251. }
  252. }