Make FunctionLibraryDefinition thread-safe.

The eager runtime mutates the FunctionLibraryRuntime's FunctionLibraryDefinition, which is shared across threads; at the same time, OpKernels might read the FunctionLibraryDefinition. This is not thread-safe unless FunctionLibraryDefinition is thread-safe.

This change makes FunctionLibraryDefinition, which is basically a map from function names to FunctionDefs, thread-safe. This is almost entirely accomplished by guarding the map with a mutex. There is however one complication: Find and RemoveFunction cannot be made thread-safe in a straightforward way (Find returns a raw pointer to a FunctionDef while Remove can delete the corresponding FunctionDef). In light of the fact that clients only ever call RemoveFunction when they in fact want to replace an existing function with a new one, we make the following modifications to FunctionLibraryDefinition's API:

1. A Contains method is added to check for the existence of a function.
2. A ReplaceFunction method is added.
3. RemoveFunction and RemoveGradient are made private.

We also update clients of the FunctionLibraryDefinition to use Contains & ReplaceFunction instead of Find and RemoveFunction.

PiperOrigin-RevId: 208271076

0 files changed, 0 insertions, 0 deletions