USDScene : Cache the reading of ShaderNetworks

Typically many prims will all reference the same UsdShadeMaterial, and we want to avoid loading the same material over and over again. This change gives a 10x speedup when loading the Moana island scene in Gaffer.

Author: johnhaddon

contrib/IECoreUSD/src/IECoreUSD/USDScene.cpp

@@ -41,6 +41,7 @@
 
 #include "IECoreScene/ShaderNetwork.h"
 
+#include "IECore/LRUCache.h"
 #include "IECore/MessageHandler.h"
 #include "IECore/SimpleTypedData.h"
 #include "IECore/VectorTypedData.h"
@@ -389,6 +390,62 @@ void populateMaterial( pxr::UsdShadeMaterial &mat, const std::map< const Interne
 	}
 }
 
+/// SdfPath is the appropriate cache key for _storage_, but we need a
+/// `UsdShadeOutput` for computation. This struct provides the implicit
+/// conversion that LRUCache needs to make that possible.
+struct ShaderNetworkCacheGetterKey : public pxr::UsdShadeOutput
+{
+	ShaderNetworkCacheGetterKey( const pxr::UsdShadeOutput &output )
+		:	pxr::UsdShadeOutput( output )
+	{
+	}
+
+	operator pxr::SdfPath () const
+	{
+		return GetAttr().GetPath();
+	}
+};
+
+class ShaderNetworkCache : public LRUCache<pxr::SdfPath, IECoreScene::ConstShaderNetworkPtr, LRUCachePolicy::Parallel, ShaderNetworkCacheGetterKey>
+{
+
+	public :
+
+		ShaderNetworkCache( size_t maxBytes )
+			:	LRUCache<pxr::SdfPath, IECoreScene::ConstShaderNetworkPtr, LRUCachePolicy::Parallel, ShaderNetworkCacheGetterKey>( getter, maxBytes )
+		{
+		}
+
+	private :
+
+		static IECoreScene::ConstShaderNetworkPtr getter( const ShaderNetworkCacheGetterKey &key, size_t &cost )
+		{
+			IECoreScene::ConstShaderNetworkPtr result;
+
+			/// \todo I'm pretty sure that the `readShaderNetwork()` signature is overly complex,
+			/// and it should just be passed a single `UsdShadeOutput &` like this function.
+			/// I suspect that `writeShaderNetwork()` could take a single `UsdShadeOutput &` too,
+			/// for symmetry between the two functions.
+
+			pxr::UsdShadeConnectableAPI source;
+			pxr::TfToken sourceName;
+			pxr::UsdShadeAttributeType sourceType;
+			if( key.GetConnectedSource( &source, &sourceName, &sourceType ) )
+			{
+				pxr::UsdShadeShader s( source.GetPrim() );
+				result = ShaderAlgo::readShaderNetwork( source.GetPrim().GetParent().GetPath(), s, sourceName );
+			}
+			else
+			{
+				result = new IECoreScene::ShaderNetwork();
+			}
+
+			cost = result->Object::memoryUsage();
+			return result;
+		}
+
+};
+
 } // namespace
 
 class USDScene::Location : public RefCounted
@@ -411,7 +468,8 @@ class USDScene::IO : public RefCounted
 		IO( const std::string &fileName, const pxr::UsdStageRefPtr &stage, IndexedIO::OpenMode openMode )
 			:	m_fileName( fileName ), m_openMode( openMode ), m_stage( stage ),
 				m_rootPrim( m_stage->GetPseudoRoot() ),
-				m_timeCodesPerSecond( m_stage->GetTimeCodesPerSecond() )
+				m_timeCodesPerSecond( m_stage->GetTimeCodesPerSecond() ),
+				m_shaderNetworkCache( 10 * 1024 * 1024 ) // 10Mb
 		{
 		}
 
@@ -489,6 +547,11 @@ class USDScene::IO : public RefCounted
 			);
 		}
 
+		IECoreScene::ConstShaderNetworkPtr readShaderNetwork( const pxr::UsdShadeOutput &output )
+		{
+			return m_shaderNetworkCache.get( output );
+		}
+
 	private :
 
 		static pxr::UsdStageRefPtr makeStage( const std::string &fileName, IndexedIO::OpenMode openMode )
@@ -522,6 +585,8 @@ class USDScene::IO : public RefCounted
 		pxr::UsdShadeMaterialBindingAPI::BindingsCache m_usdBindingsCache;
 		pxr::UsdShadeMaterialBindingAPI::CollectionQueryCache m_usdCollectionQueryCache;
 
+		ShaderNetworkCache m_shaderNetworkCache;
+
 };
 
 USDScene::USDScene( const std::string &fileName, IndexedIO::OpenMode openMode )
@@ -910,17 +975,7 @@ ConstObjectPtr USDScene::readAttribute( const SceneInterface::Name &name, double
 			pxr::UsdShadeOutput o = mat.GetOutput( n );
 			if( o && pxr::UsdAttribute( o ).IsAuthored() )
 			{
-				{
-					pxr::UsdShadeConnectableAPI source;
-					pxr::TfToken sourceName;
-					pxr::UsdShadeAttributeType sourceType;
-					if( o.GetConnectedSource( &source, &sourceName, &sourceType ) )
-					{
-						pxr::UsdShadeShader s( source.GetPrim() );
-						return ShaderAlgo::readShaderNetwork( source.GetPrim().GetParent().GetPath(), s, sourceName );
-					}
-				}
-				return new IECoreScene::ShaderNetwork();
+				return m_root->readShaderNetwork( o );
 			}
 		}
 	}