[PyTorch Edge] Don't use LeftRight in mobile (#66064)

Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/66064

The only place this is used seems to be in the dispatcher for `operatorLookupTable_`. Disarming `LeftRight` disarms it for this one use case.

This should make .so loading faster, and also reduce memory consumption since `LeftRight<T>` does 2 writes for every write. I'd like to get a thorough review from reviewers for this diff since I want to make sure that initialization of stuff that writes into the dispatcher isn't going to happen on multiple threads for on-device use.

Created a new class named `LeftRightNoOpWrapper<T>` for use in mobile builds.

### Why is LeftRight<T> slow?

It maintains 2 copies of each data structure `T` to be able to keep reads quick. Every write goes to both data structures, which means that writes that 2x and memory overhead is also 2x

### Why is this safe for mobile builds?

1. .so loading never happens concurrently with model execution
2. Custom ops are loaded during .so load - initializers are all run serially
3. I don't see any threads being spawned from the global schema and kernel initializers

After discussing with dreiss, it seems like there could be rare cases in OSS apps or internal Android/iOS apps where a `.so` or `dylib` is loaded after the PT runtime is loaded, and this load happens concurrently with an in-progress inference run, which is looking up the operator table in the dispatcher.

To avoid crashes there, it seems reasonable to use the RW lock, since I don't expect any contention 99.9% of the time.

When registering operators, everything is serial so only one thread will ever hold the lock. The next time it needs the lock, it will have already released it.
During inference runs, only one thread will ask for the shared lock unless multiple concurrent inferences are in progress. Even in that case, they will all be able to simultaneously get the Read lock.

Test Plan: Build and generate a local build of the iOS app to test.

Reviewed By: swolchok

Differential Revision: D31352346

fbshipit-source-id: c3f12454de3dbd7b421a6057d561e9373ef5bf98

aten/src/ATen/core/dispatch/Dispatcher.h

@@ -283,7 +283,11 @@ private:
   void checkSchemaCompatibility(const OperatorHandle& op, const FunctionSchema& schema, const std::string& debug);
 
   std::list<OperatorDef> operators_;
+#if !defined(C10_MOBILE)
   LeftRight<ska::flat_hash_map<OperatorName, OperatorHandle>> operatorLookupTable_;
+#else
+  RWSafeLeftRightWrapper<ska::flat_hash_map<OperatorName, OperatorHandle>> operatorLookupTable_;
+#endif
   // Map from namespace to debug string (saying, e.g., where the library was defined)
   ska::flat_hash_map<std::string, std::string> libraries_;
 

c10/util/LeftRight.h

@@ -3,6 +3,7 @@
 #include <atomic>
 #include <functional>
 #include <mutex>
+#include <shared_mutex>
 #include <thread>
 
 namespace c10 {
@@ -187,4 +188,40 @@ class LeftRight final {
   std::mutex _writeMutex;
 };
 
+// RWSafeLeftRightWrapper is API compatible with LeftRight and uses a
+// read-write lock to protect T (data).
+template <class T>
+class RWSafeLeftRightWrapper final {
+  using mutexType = std::mutex;
+  using rLockType = std::unique_lock<std::mutex>;
+  using wLockType = std::unique_lock<std::mutex>;
+
+ public:
+  template <class... Args>
+  explicit RWSafeLeftRightWrapper(const Args&... args) : _data{args...} {}
+
+  // RWSafeLeftRightWrapper is not copyable or moveable since LeftRight
+  // is not copyable or moveable.
+  RWSafeLeftRightWrapper(const RWSafeLeftRightWrapper&) = delete;
+  RWSafeLeftRightWrapper(RWSafeLeftRightWrapper&&) noexcept = delete;
+  RWSafeLeftRightWrapper& operator=(const RWSafeLeftRightWrapper&) = delete;
+  RWSafeLeftRightWrapper& operator=(RWSafeLeftRightWrapper&&) noexcept = delete;
+
+  template <typename F>
+  auto read(F&& readFunc) const -> typename std::result_of<F(const T&)>::type {
+    rLockType lock(mutex_);
+    return readFunc(_data);
+  }
+
+  template <typename F>
+  auto write(F&& writeFunc) -> typename std::result_of<F(T&)>::type {
+    wLockType lock(mutex_);
+    return writeFunc(_data);
+  }
+
+ private:
+  T _data;
+  mutable mutexType mutex_;
+};
+
 } // namespace c10