[mlir] Refactor InterfaceMap to use a sorted vector of interfaces, as opposed to a DenseMap

[rriddle]

Summary

A majority of operations have a very small number of interfaces, which means that the cost of using a hash map is generally larger for interface lookups than just a binary search. In the future when there are a number of operations with large amounts of interfaces, we can switch to a hybrid approach that optimizes lookups based on the number of interfaces. For now, however, a binary search is the best approach.

This dropped compile time on a largish TF MLIR module by 20%(half a second).

[jpienaar]

Nice

mlir/include/mlir/IR/OperationSupport.h

330–332

Was this change needed too?

mlir/include/mlir/Support/InterfaceSupport.h

195

Could this be hoisted into a helper function that we could use both places?

[rriddle]

mlir/include/mlir/IR/OperationSupport.h

330–332

Yes, it's important for ensuring that interface lookups get optimized properly. In most cases, this function gets inlined.

mlir/include/mlir/IR/OperationSupport.h

public:
   /// Return the name of this operation. This always succeeds.
   StringRef getStringRef() const;
 
   /// Return the name of this operation as an identifier. This always succeeds.
   Identifier getIdentifier() const;
 
   /// If this operation has a registered operation description, return it.
   /// Otherwise return null.
-  const AbstractOperation *getAbstractOperation() const;
+  const AbstractOperation *getAbstractOperation() const {
+    return representation.dyn_cast<const AbstractOperation *>();
+  }
 
   void print(raw_ostream &os) const;
   void dump() const;
 
   void *getAsOpaquePointer() const {
     return static_cast<void *>(representation.getOpaqueValue());
   }
   static OperationName getFromOpaquePointer(const void *pointer);

mlir/include/mlir/Support/InterfaceSupport.h

class InterfaceMap {
   using detect_get_interface_id = llvm::is_detected<has_get_interface_id, T>;
   template <typename... Types>
   using num_interface_types = typename std::tuple_size<
       typename FilterTypes<detect_get_interface_id, Types...>::type>;
 
 public:
   InterfaceMap(InterfaceMap &&) = default;
   ~InterfaceMap() {
-    if (interfaces) {
-      for (auto &it : *interfaces)
-        free(it.second);
-    }
-  }
+    for (auto &it : interfaces)
 
   /// Construct an InterfaceMap with the given set of template types. For
   /// convenience given that object trait lists may contain other non-interface
   /// types, not all of the types need to be interfaces. The provided types that
   /// do not represent interfaces are not added to the interface map.
   template <typename... Types>
   static std::enable_if_t<num_interface_types<Types...>::value != 0,
                           InterfaceMap>
static std::enable_if_t<num_interface_types<Types...>::value == 0,
                           InterfaceMap>
   get() {
     return InterfaceMap();
   }
 
   /// Returns an instance of the concept object for the given interface if it
   /// was registered to this map, null otherwise.
   template <typename T> typename T::Concept *lookup() const {
-    void *inst = interfaces ? interfaces->lookup(T::getInterfaceID()) : nullptr;
-    return reinterpret_cast<typename T::Concept *>(inst);
+    return reinterpret_cast<typename T::Concept *>(lookup(T::getInterfaceID()));
   }
 
 private:
+  /// Compare two TypeID instances by comparing the underlying pointer.
+  static bool compare(TypeID lhs, TypeID rhs) {
+    return lhs.getAsOpaquePointer() < rhs.getAsOpaquePointer();
+  }
+
   InterfaceMap() = default;
   InterfaceMap(MutableArrayRef<std::pair<TypeID, void *>> elements)
-      : interfaces(std::make_unique<llvm::SmallDenseMap<TypeID, void *>>(
-            elements.begin(), elements.end())) {}
+      : interfaces(elements.begin(), elements.end()) {
+    llvm::sort(interfaces, [](const auto &lhs, const auto &rhs) {
+      return compare(lhs.first, rhs.first);
+    });
+  }
 
   template <typename... Ts>
   static InterfaceMap getImpl(std::tuple<Ts...> *) {
     std::pair<TypeID, void *> elements[] = {std::make_pair(
         Ts::getInterfaceID(),
         new (malloc(sizeof(typename Ts::ModelT))) typename Ts::ModelT())...};
     return InterfaceMap(elements);
   }
 
-  /// The internal map of interfaces. This is constructed statically for each
-  /// set of interfaces.
-  std::unique_ptr<llvm::SmallDenseMap<TypeID, void *>> interfaces;
+  /// Returns an instance of the concept object for the given interface id if it
+  /// was registered to this map, null otherwise.
+  void *lookup(TypeID id) const {
+    auto it = llvm::lower_bound(interfaces, id, [](const auto &it, TypeID id) {
+      return compare(it.first, id);
+    });
+    return (it != interfaces.end() && it->first == id) ? it->second : nullptr;
+  }
+
+  /// A list of interface instances, sorted by TypeID.
+  SmallVector<std::pair<TypeID, void *>> interfaces;
 };
 
 } // end namespace detail
 } // end namespace mlir
 
 #endif

mlir/lib/IR/Operation.cpp

 
 /// Return the name of this operation as an identifier. This always succeeds.
 Identifier OperationName::getIdentifier() const {
   if (auto *op = representation.dyn_cast<const AbstractOperation *>())
     return op->name;
   return representation.get<Identifier>();
 }
 
-const AbstractOperation *OperationName::getAbstractOperation() const {
-  return representation.dyn_cast<const AbstractOperation *>();
-}
-
 OperationName OperationName::getFromOpaquePointer(const void *pointer) {
   return OperationName(
       RepresentationUnion::getFromOpaqueValue(const_cast<void *>(pointer)));
 }
 
 //===----------------------------------------------------------------------===//
 // Operation
 //===----------------------------------------------------------------------===//