[PyTorch][JIT] use a better hash table in alias analysis

Pull Request resolved: #69854


ghstack-source-id: 145889449

Differential Revision: [D33039733](https://our.internmc.facebook.com/intern/diff/D33039733/)

Author: swolchok

torch/csrc/jit/ir/alias_analysis.cpp

@@ -1,5 +1,6 @@
 #include <torch/csrc/jit/ir/alias_analysis.h>
 
+#include <c10/util/flat_hash_map.h>
 #include <c10/util/irange.h>
 #include <torch/csrc/jit/jit_log.h>
 #include <torch/csrc/jit/passes/utils/subgraph_utils.h>
@@ -26,7 +27,7 @@ c10::MaybeOwned<TypePtr> toSingleType(const AliasTypeSet& mut_types) {
 class MutableTypePtrHelper {
  public:
   explicit MutableTypePtrHelper(
-      std::unordered_map<TypePtr, AliasTypeSet>* mutable_type_cache)
+      ska::flat_hash_map<TypePtr, AliasTypeSet>* mutable_type_cache)
       : mutable_type_cache_(mutable_type_cache) {}
 
   // Map any mutable type to a type such that all other types which the
@@ -140,12 +141,12 @@ class MutableTypePtrHelper {
         return c10::nullopt;
     }
   }
-  std::unordered_map<TypePtr, AliasTypeSet>* mutable_type_cache_;
+  ska::flat_hash_map<TypePtr, AliasTypeSet>* mutable_type_cache_;
 };
 
 bool isMutableTypeImpl(
     const TypePtr& type,
-    std::unordered_map<TypePtr, AliasTypeSet>* mutable_type_cache) {
+    ska::flat_hash_map<TypePtr, AliasTypeSet>* mutable_type_cache) {
   // Check common cases to avoid recursively constructing type in
   // `mapTypeToAliasTypeSetPtrImpl`
   auto kind = type->kind();

torch/csrc/jit/ir/alias_analysis.h

@@ -260,7 +260,7 @@ class AliasDb {
   // Mapping of values to MemoryDAG elements
   ska::flat_hash_map<const Value*, Element*> elementMap_;
   // All wildcard Elements (one for each unique mutable type)
-  std::unordered_map<TypePtr, Element*, HashType, EqualType> wildcardIndex_;
+  ska::flat_hash_map<TypePtr, Element*, HashType, EqualType> wildcardIndex_;
   Element* getWildcard(const TypePtr& type) const;
   c10::optional<Element*> tryGetOrCreateWildcard(const TypePtr& type);
   void addContainedTypesToFreshElement(
@@ -276,7 +276,7 @@ class AliasDb {
   bool hasWriters(const at::ArrayRef<Value*>& values) const;
 
   // Cached mapping of type ptrs to their mutable types
-  mutable std::unordered_map<TypePtr, AliasTypeSet> mapped_mutable_types_;
+  mutable ska::flat_hash_map<TypePtr, AliasTypeSet> mapped_mutable_types_;
 
   /**
    * State for tracking write info.