Instancer : Parallelize prototype creation when encapsulated

src/GafferScene/Instancer.cpp

@@ -2650,39 +2650,116 @@ void Instancer::InstancerCapsule::render( IECoreScenePreview::Renderer *renderer
 
 	if( engines[0]->hasContextVariables() )
 	{
+		// Since every instance could have a different prototype could be different depending on the context,
+		// we need to loop through every instance, set up the context so we can get a hash from it, and if
+		// we haven't yet set up a prototype for that context, set one up.
+		//
+		// This could expensive if there are many distinct prototypes, so this is a bit complicated in order
+		// to parallelize it.
+
+		// Store the index in the `prototypes` list corresponding to each context hash
 		std::unordered_map< IECore::MurmurHash, int, MurmurHashHasher > prototypeContextsMap;
-		prototypeIndices.reserve( pointIndicesForPrototype.size() );
 
-		Context::EditableScope prototypeScope( scope.context() );
+		// Each point gets an index telling it which prototype to use in the `prototypes` vector
+		prototypeIndices.resize( pointIndicesForPrototype.size() );
 
-		for( int pointIndex : pointIndicesForPrototype )
-		{
+		// Must be held to read prototypeContextsMap, must have write lock to modify it.
+		tbb::spin_rw_mutex prototypeContextsMapMutex;
 
-			// We find the capsules using the engine at shutter open, but the time used to construct the capsules
-			// must be the on-frame time, since the capsules will add their own shutter ( and we also handle
-			// the shutter ourselves for transform matrices )
-			//
-			// For most context variables, we are overwriting them for each prototype anyway, so
-			// we can reuse the context. But timeOffset is relative, so it's important that we reset the
-			// time before we do setPrototypeContextVariables for the next element. ( Should this be more
-			// general instead of assuming that frame is the only variable for which offsetMode may be set? )
-			prototypeScope.setFrame( onFrameTime );
+		// Must be held to read or modify elements of the `prototypes` vector.
+		// A write lock is only required to resize the vector.
+		tbb::spin_rw_mutex prototypesMutex;
 
-			engines[0]->setPrototypeContextVariables( pointIndex, prototypeScope );
-			auto inserted = prototypeContextsMap.emplace( prototypeScope.context()->hash(), prototypeContextsMap.size() );
-
-			if( inserted.second )
+		task_group_context taskGroupContext( task_group_context::isolated );
+		tbb::parallel_for( tbb::blocked_range<size_t>( 0, pointIndicesForPrototype.size() ), [&]( const tbb::blocked_range<size_t> &r )
 			{
-				prototypes.push_back( 
-					addPrototypeToVectors(
-						prototypesPlug, prototypeRoot, prototypeScope, sampleTimes, capsuleHash, *renderOptions
-					)
-				);
-			}
+				Context::EditableScope prototypeScope( scope.context() );
+				for( size_t i = r.begin(); i != r.end(); i++ )
+				{
+					int pointIndex = pointIndicesForPrototype[i];
+
+					// We find the capsules using the engine at shutter open, but the time used to construct the capsules
+					// must be the on-frame time, since the capsules will add their own shutter ( and we also handle
+					// the shutter ourselves for transform matrices )
+					//
+					// For most context variables, we are overwriting them for each prototype anyway, so
+					// we can reuse the context. But timeOffset is relative, so it's important that we reset the
+					// time before we do setPrototypeContextVariables for the next element. ( Should this be more
+					// general instead of assuming that frame is the only variable for which offsetMode may be set? )
+					prototypeScope.setFrame( onFrameTime );
+
+					engines[0]->setPrototypeContextVariables( pointIndex, prototypeScope );
+
+					MurmurHash contextHash = prototypeScope.context()->hash();
+
+					// We only need a read lock to check if there is already a prototype corresponding to this
+					// context.
+					tbb::spin_rw_mutex::scoped_lock prototypeContextsLock( prototypeContextsMapMutex, false );
+					size_t prototypeIndex;
+					auto existingPrototypeIndex = prototypeContextsMap.find( contextHash );
+					if( existingPrototypeIndex != prototypeContextsMap.end() )
+					{
+						// Found an existing prototype, we just use this index
+						prototypeIndex = existingPrototypeIndex->second;
+						prototypeContextsLock.release();
+					}
+					else
+					{
+						// Need to add a new prototype, so we need a write lock on the contexts map to
+						// add a new index
+						prototypeContextsLock.upgrade_to_writer();
+						auto inserted = prototypeContextsMap.emplace( contextHash, prototypeContextsMap.size() );
 
-			prototypeIndices.push_back( inserted.first->second );
+						prototypeIndex = inserted.first->second;
 
-		}
+						// We now have an index, and no longer need to access the contexts map.
+						prototypeContextsLock.release();
+
+						if( !inserted.second )
+						{
+							// Someone else allocated a prototype for this context while we were acquiring the
+							// write lock, we don't need to add anything after all
+						}
+						else
+						{
+							// We're actually adding a prototype, so we need the `prototypesLock`
+							tbb::spin_rw_mutex::scoped_lock prototypesLock( prototypesMutex, false );
+
+							if( prototypes.size() <= prototypeIndex )
+							{
+								// We need to resize the `prototypes` vector, this is the only time
+								// we need a write lock.
+								prototypesLock.upgrade_to_writer();
+
+								// Double check that someone else hasn't resized it for us while we were getting
+								// the lock
+								if( prototypes.size() <= prototypeIndex )
+								{
+									// Resize up by a factor of 2, to amortize how many resizes we need to do
+									prototypes.resize( ( prototypeIndex + 1 ) * 2 );
+								}
+								prototypesLock.downgrade_to_reader();
+							}
+
+							// We've gotten a unique prototypeIndex from the `prototypeContextsMap`, and each
+							// element of `prototypes` can be set independently, so we can safely make this
+							// assignment with just a read lock on prototypesMutex
+
+							prototypes[prototypeIndex] = addPrototypeToVectors(
+									prototypesPlug, prototypeRoot, prototypeScope, sampleTimes, capsuleHash, *renderOptions
+							);
+						}
+					}
+
+					prototypeIndices[i] = prototypeIndex;
+				}
+			},
+			taskGroupContext
+		);
+
+		// The prototypes vector is allocated over-sized, so we don't need to reallocate it for each prototype we
+		// add. The actual number of prototypes can be found from the size of the prototypesContextsMap.
+		prototypes.resize( prototypeContextsMap.size() );
 	}
 	else
 	{