Add OpenExec support for caching simple prim definitions, by introducing a

"composed prim definition" structure that stores computations defined on all
ancestor types of the provider prim's schema type.

(Internal change: 2366612)

Author: adamrwoodbury

pxr/exec/exec/definitionRegistry.cpp

@@ -124,35 +124,24 @@ Exec_DefinitionRegistry::GetComputationDefinition(
         return nullptr;
     }
 
-    // Iterate over all ancestor types of the provider's schema type, from
-    // derived to base, starting with the schema type itself. Look for a
-    // matching plugin prim computation for the derived-most schema type that
-    // defines it, or null, if no matching computation can be found.
-    //
-    // TODO: Repeatedly traversing the schema type hierarchy like this is
-    // wasteful and we plan to cache results appropriately. But we still need to
-    // add support for applied schemas, and we will also need to add the ability
-    // to compose computation definitions, so for now we are leaving this
-    // inefficiency in place until more of that functionality lands. The current
-    // thinking is that exec will cache composed prim definitions that will be
-    // keyed off of a tuple of typed and applied schemas and will cache the
-    // resulting set of composed computation definitions. That will enable this
-    // code to construct a key and do a single lookup into the cache, rather
-    // than searching to find the computation definition.
-
-    TfType foundType;
-    std::vector<TfType> schemaAncestorTypes;
-    schemaType.GetAllAncestorTypes(&schemaAncestorTypes);
-
-    for (const TfType type : schemaAncestorTypes) {
-        if (const auto pluginIt = _pluginPrimComputationDefinitions.find(
-                {type, computationName});
-            pluginIt != _pluginPrimComputationDefinitions.end()) {
-            return &pluginIt->second;
-        }
+    // Get the composed prim definition, creating it if necesseary, and use it
+    // to look up the computation, or to determine that the requested
+    // computation isn't defined for this prim.
+    auto composedDefIt = _composedPrimDefinitions.find(schemaType);
+    if (composedDefIt == _composedPrimDefinitions.end()) {
+        // Note that we allow concurrent callers to race to compose prim
+        // definitions, since it is safe to do so and we don't expect it to
+        // happen in the common case.
+        _ComposedPrimDefinition primDef =
+            _ComposePrimDefinition(schemaType);
+
+        composedDefIt = _composedPrimDefinitions.emplace(
+            schemaType, std::move(primDef)).first;
     }
 
-    return nullptr;
+    const auto &compDefs = composedDefIt->second.primComputationDefinitions;
+    const auto it = compDefs.find(computationName);
+    return it == compDefs.end() ? nullptr : it->second;
 }
 
 const Exec_ComputationDefinition *
@@ -176,6 +165,52 @@ Exec_DefinitionRegistry::GetComputationDefinition(
     return nullptr;
 }
 
+Exec_DefinitionRegistry::_ComposedPrimDefinition
+Exec_DefinitionRegistry::_ComposePrimDefinition(
+    const TfType schemaType) const
+{
+    TRACE_FUNCTION();
+
+    // Iterate over all ancestor types of the provider's schema type, from
+    // derived to base, starting with the schema type itself. Ensure that plugin
+    // computations have been loaded for each schema type for which they are
+    // registered. Add all plugin computations registered for each type to the
+    // composed prim definition.
+    //
+    // TODO: Add support for computations that are registered for applied
+    // schemas. To do that, instead of keying off the schema type we will need
+    // to use a "configuration key" that combines the typed schema with applied
+    // schemas. We will also need to search through all applied schemas, in
+    // strength order, in addition to searching up the typed schema type
+    // hierarchy.
+
+    std::vector<TfType> schemaAncestorTypes;
+    schemaType.GetAllAncestorTypes(&schemaAncestorTypes);
+
+    // Build up the composed prim definition.
+    _ComposedPrimDefinition primDef;
+
+    for (const TfType type : schemaAncestorTypes) {
+
+        // TODO: For all but the first type, it makes sense to look in
+        // _composedPrimDefinitions to see if we have already composed the base
+        // type, and then to merge, rather than keep searching up the type
+        // hierarchy.
+
+        if (const auto pluginIt = _pluginPrimComputationDefinitions.find(type);
+            pluginIt != _pluginPrimComputationDefinitions.end()) {
+            for (const Exec_PluginComputationDefinition &computationDef :
+                     pluginIt->second) {
+                primDef.primComputationDefinitions.emplace(
+                    computationDef.GetComputationName(),
+                    &computationDef);
+            }
+        }
+    }
+
+    return primDef;
+}
+
 void
 Exec_DefinitionRegistry::_RegisterPrimComputation(
     TfType schemaType,
@@ -203,14 +238,11 @@ Exec_DefinitionRegistry::_RegisterPrimComputation(
     }
 
     const bool emplaced =
-        _pluginPrimComputationDefinitions.emplace(
-            std::piecewise_construct, 
-            std::forward_as_tuple(schemaType, computationName),
-            std::forward_as_tuple(
-                resultType,
-                computationName,
-                std::move(callback),
-                std::move(inputKeys))).second;
+        _pluginPrimComputationDefinitions[schemaType].emplace(
+            resultType,
+            computationName,
+            std::move(callback),
+            std::move(inputKeys)).second;
 
     if (!emplaced) {
         TF_CODING_ERROR(

pxr/exec/exec/definitionRegistry.h

@@ -20,7 +20,10 @@
 #include "pxr/base/tf/type.h"
 #include "pxr/base/tf/weakBase.h"
 
+#include "tbb/concurrent_unordered_map.h"
+
 #include <memory>
+#include <set>
 #include <tuple>
 #include <unordered_map>
 #include <utility>
@@ -97,10 +100,32 @@ class Exec_DefinitionRegistry : public TfWeakBase
     Exec_DefinitionRegistry();
 
     // Returns a reference to the singleton that is suitable for registering
-    // new computations. This instance cannot be used to look up computations.
+    // new computations.
+    //
+    // The returned instance cannot be used to look up computations.
+    //
     EXEC_API
     static Exec_DefinitionRegistry& _GetInstanceForRegistration();
 
+    // a structure that contains the definitions for all computations that can
+    // be found on a prim of a given type.
+    //
+    struct _ComposedPrimDefinition {
+        // Map from computation name to plugin prim computation definition.
+        std::unordered_map<
+            TfToken,
+            const Exec_PluginComputationDefinition *,
+            TfHash>
+        primComputationDefinitions;
+
+        // TODO: Add plugin attribute computation definitions.
+    };
+
+    // Creates and returns the composed prim definition for a prim with type
+    // \p schemaType.
+    //
+    _ComposedPrimDefinition _ComposePrimDefinition(TfType schemaType) const;
+
     void _RegisterPrimComputation(
         TfType schemaType,
         const TfToken &computationName,
@@ -130,17 +155,34 @@ class Exec_DefinitionRegistry : public TfWeakBase
     // functions.
     std::unique_ptr<Exec_RegistrationBarrier> _registrationBarrier;
 
-    // Map from (schemaType, computationName) to plugin prim computation
-    // definition.
-    //
-    // TODO: When we add support for loading plugins, we will need to think about
-    // how to provide thread-safe access to this map.
+    // Comparator for ordering plugin computation definitions in a set.
+    struct _PluginComputationDefinitionComparator {
+        bool operator()(
+            const Exec_PluginComputationDefinition &a,
+            const Exec_PluginComputationDefinition &b) const {
+            return a.GetComputationName() < b.GetComputationName();
+        }
+    };
+
+    // Map from schemaType to plugin prim computation definitions.
     std::unordered_map<
-        std::tuple<TfType, TfToken>,
-        Exec_PluginComputationDefinition,
+        TfType,
+        std::set<
+            Exec_PluginComputationDefinition,
+            _PluginComputationDefinitionComparator>,
         TfHash>
     _pluginPrimComputationDefinitions;
 
+    // Map from schemaType to composed prim exec definition.
+    //
+    // This is a concurrent map to allow computation lookup to happen in
+    // parallel with lazy caching of composed prim definitions.
+    mutable tbb::concurrent_unordered_map<
+        TfType,
+        _ComposedPrimDefinition,
+        TfHash>
+    _composedPrimDefinitions;
+
     // Map from computationName to builtin stage computation
     // definition.
     std::unordered_map<

pxr/exec/exec/testenv/testExecComputationRegistration.cpp

@@ -40,6 +40,7 @@ TF_DEFINE_PRIVATE_TOKENS(
     (attributeComputation)
     (attributeComputedValueComputation)
     (attributeName)
+    (baseAndDerivedSchemaComputation)
     (derivedSchemaComputation)
     (emptyComputation)
     (missingComputation)
@@ -105,12 +106,23 @@ EXEC_REGISTER_SCHEMA(TestExecComputationRegistrationCustomSchema)
             Attribute(_tokens->attr)
                 .Computation<double>(ExecBuiltinComputations->computeValue)
         );
+
+    self.PrimComputation(_tokens->baseAndDerivedSchemaComputation)
+        .Callback(+[](const VdfContext &) { return 1.0; });
 }
 
 EXEC_REGISTER_SCHEMA(TestExecComputationRegistrationDerivedCustomSchema)
 {
     self.PrimComputation(_tokens->derivedSchemaComputation)
         .Callback(+[](const VdfContext &) { return 1.0; });
+
+    // This overrides the computation of the same name on the base schema.
+    // (We add an input here so we can verify this definition is stronger.)
+    self.PrimComputation(_tokens->baseAndDerivedSchemaComputation)
+        .Callback(+[](const VdfContext &) { return 1.0; })
+        .Inputs(
+            AttributeValue<int>(_tokens->attributeName)
+        );
 }
 
 // XXX:TODO
@@ -461,6 +473,21 @@ TestDerivedSchemaComputationRegistration()
         TF_AXIOM(primCompDef);
     }
 
+    {
+        // Look up a computation registered for the base and derived schema
+        // types.
+        const Exec_ComputationDefinition *const primCompDef =
+            reg.GetComputationDefinition(
+                *prim, _tokens->baseAndDerivedSchemaComputation, nullJournal);
+        TF_AXIOM(primCompDef);
+
+        // Make sure we got the definition from the derived schema (i.e., the
+        // stronger one).
+        const auto inputKeys =
+            primCompDef->GetInputKeys(*prim, nullJournal);
+        ASSERT_EQ(inputKeys.size(), 1);
+    }
+
     {
         // Look up a computation registered for the base schema type.
         const Exec_ComputationDefinition *const primCompDef =