[mlir][StorageUniquer] Replace all usages of std::function with function_ref.

[rriddle]

Summary

std::function has a notoriously large amount of malloc traffic, whereas function_ref is a cheaper and more efficient alternative.

[lattner]

awesome, thank you!

mlir/include/mlir/IR/AttributeSupport.h

 // A utility class to get, or create, unique instances of attributes within an
 // MLIRContext. This class manages all creation and uniquing of attributes.
 class AttributeUniquer {
 public:
   /// Get an uniqued instance of attribute T.
   template <typename T, typename... Args>
   static T get(MLIRContext *ctx, unsigned kind, Args &&... args) {
     return ctx->getAttributeUniquer().get<typename T::ImplType>(
-        getInitFn(ctx, T::getTypeID()), kind, std::forward<Args>(args)...);
+        [ctx](AttributeStorage *storage) {
+          initializeAttributeStorage(storage, ctx, T::getTypeID());
+        },
+        kind, std::forward<Args>(args)...);
   }
 
 private:
-  /// Returns a functor used to initialize new attribute storage instances.
-  static std::function<void(AttributeStorage *)> getInitFn(MLIRContext *ctx,
-                                                           TypeID attrID);
+  /// Initialize the given attribute storage instance.
+  static void initializeAttributeStorage(AttributeStorage *storage,
+                                         MLIRContext *ctx, TypeID attrID);
 };
 } // namespace detail
 
 } // end namespace mlir
 
 #endif

mlir/include/mlir/Support/StorageUniquer.h

private:
     llvm::BumpPtrAllocator allocator;
   };
 
   /// Gets a uniqued instance of 'Storage'. 'initFn' is an optional parameter
   /// that can be used to initialize a newly inserted storage instance. This
   /// function is used for derived types that have complex storage or uniquing
   /// constraints.
   template <typename Storage, typename Arg, typename... Args>
-  Storage *get(std::function<void(Storage *)> initFn, unsigned kind, Arg &&arg,
+  Storage *get(function_ref<void(Storage *)> initFn, unsigned kind, Arg &&arg,
                Args &&... args) {
     // Construct a value of the derived key type.
     auto derivedKey =
         getKey<Storage>(std::forward<Arg>(arg), std::forward<Args>(args)...);
 
     // Create a hash of the kind and the derived key.
     unsigned hashValue = getHash<Storage>(kind, derivedKey);
 
     // Generate an equality function for the derived storage.
-    std::function<bool(const BaseStorage *)> isEqual =
-        [&derivedKey](const BaseStorage *existing) {
-          return static_cast<const Storage &>(*existing) == derivedKey;
-        };
+    auto isEqual = [&derivedKey](const BaseStorage *existing) {
 
     // Generate a constructor function for the derived storage.
-    std::function<BaseStorage *(StorageAllocator &)> ctorFn =
-        [&](StorageAllocator &allocator) {
-          auto *storage = Storage::construct(allocator, derivedKey);
-          if (initFn)
-            initFn(storage);
-          return storage;
-        };
+    auto ctorFn = [&](StorageAllocator &allocator) {
 
     // Get an instance for the derived storage.
     return static_cast<Storage *>(getImpl(kind, hashValue, isEqual, ctorFn));
   }
 
   /// Gets a uniqued instance of 'Storage'. 'initFn' is an optional parameter
   /// that can be used to initialize a newly inserted storage instance. This
   /// function is used for derived types that use no additional storage or
   /// uniquing outside of the kind.
   template <typename Storage>
-  Storage *get(std::function<void(Storage *)> initFn, unsigned kind) {
+  Storage *get(function_ref<void(Storage *)> initFn, unsigned kind) {
     auto ctorFn = [&](StorageAllocator &allocator) {
       auto *storage = new (allocator.allocate<Storage>()) Storage();
       if (initFn)
         initFn(storage);
       return storage;
     };
     return static_cast<Storage *>(getImpl(kind, ctorFn));
   }
 
   /// Erases a uniqued instance of 'Storage'. This function is used for derived
   /// types that have complex storage or uniquing constraints.
   template <typename Storage, typename Arg, typename... Args>
   void erase(unsigned kind, Arg &&arg, Args &&... args) {
     // Construct a value of the derived key type.
     auto derivedKey =
         getKey<Storage>(std::forward<Arg>(arg), std::forward<Args>(args)...);
 
     // Create a hash of the kind and the derived key.
     unsigned hashValue = getHash<Storage>(kind, derivedKey);
 
     // Generate an equality function for the derived storage.
-    std::function<bool(const BaseStorage *)> isEqual =
-        [&derivedKey](const BaseStorage *existing) {
-          return static_cast<const Storage &>(*existing) == derivedKey;
-        };
+    auto isEqual = [&derivedKey](const BaseStorage *existing) {
 
     // Attempt to erase the storage instance.
     eraseImpl(kind, hashValue, isEqual, [](BaseStorage *storage) {
       static_cast<Storage *>(storage)->cleanup();
     });
   }
 
 private:
   /// Implementation for getting/creating an instance of a derived type with
   /// complex storage.
   BaseStorage *getImpl(unsigned kind, unsigned hashValue,
                        function_ref<bool(const BaseStorage *)> isEqual,
-                       std::function<BaseStorage *(StorageAllocator &)> ctorFn);
+                       function_ref<BaseStorage *(StorageAllocator &)> ctorFn);
 
   /// Implementation for getting/creating an instance of a derived type with
   /// default storage.
   BaseStorage *getImpl(unsigned kind,
-                       std::function<BaseStorage *(StorageAllocator &)> ctorFn);
+                       function_ref<BaseStorage *(StorageAllocator &)> ctorFn);
 
   /// Implementation for erasing an instance of a derived type with complex
   /// storage.
   void eraseImpl(unsigned kind, unsigned hashValue,
                  function_ref<bool(const BaseStorage *)> isEqual,
-                 std::function<void(BaseStorage *)> cleanupFn);
+                 function_ref<void(BaseStorage *)> cleanupFn);
 
   /// The internal implementation class.
   std::unique_ptr<detail::StorageUniquerImpl> impl;
 
   //===--------------------------------------------------------------------===//
   // Key Construction
   //===--------------------------------------------------------------------===//
 

mlir/lib/IR/MLIRContext.cpp

 //===----------------------------------------------------------------------===//
 
 /// Returns the storage uniquer used for constructing attribute storage
 /// instances. This should not be used directly.
 StorageUniquer &MLIRContext::getAttributeUniquer() {
   return getImpl().attributeUniquer;
 }
 
-/// Returns a functor used to initialize new attribute storage instances.
-std::function<void(AttributeStorage *)>
-AttributeUniquer::getInitFn(MLIRContext *ctx, TypeID attrID) {
-  return [ctx, attrID](AttributeStorage *storage) {
-    storage->initializeDialect(lookupDialectForSymbol(ctx, attrID));
+/// Initialize the given attribute storage instance.
+void AttributeUniquer::initializeAttributeStorage(AttributeStorage *storage,
+                                                  MLIRContext *ctx,
+                                                  TypeID attrID) {
 
-    // If the attribute did not provide a type, then default to NoneType.
-    if (!storage->getType())
-      storage->setType(NoneType::get(ctx));
-  };
 }
 
 BoolAttr BoolAttr::get(bool value, MLIRContext *context) {
   return value ? context->getImpl().trueAttr : context->getImpl().falseAttr;
 }
 
 UnitAttr UnitAttr::get(MLIRContext *context) {
   return context->getImpl().unitAttr;

mlir/lib/Support/StorageUniquer.cpp

 StorageUniquer::StorageUniquer() : impl(new StorageUniquerImpl()) {}
 StorageUniquer::~StorageUniquer() {}
 
 /// Implementation for getting/creating an instance of a derived type with
 /// complex storage.
 auto StorageUniquer::getImpl(
     unsigned kind, unsigned hashValue,
     function_ref<bool(const BaseStorage *)> isEqual,
-    std::function<BaseStorage *(StorageAllocator &)> ctorFn) -> BaseStorage * {
+    function_ref<BaseStorage *(StorageAllocator &)> ctorFn) -> BaseStorage * {
   return impl->getOrCreate(kind, hashValue, isEqual, ctorFn);
 }
 
 /// Implementation for getting/creating an instance of a derived type with
 /// default storage.
 auto StorageUniquer::getImpl(
-    unsigned kind, std::function<BaseStorage *(StorageAllocator &)> ctorFn)
+    unsigned kind, function_ref<BaseStorage *(StorageAllocator &)> ctorFn)
     -> BaseStorage * {
   return impl->getOrCreate(kind, ctorFn);
 }
 
 /// Implementation for erasing an instance of a derived type with complex
 /// storage.
 void StorageUniquer::eraseImpl(unsigned kind, unsigned hashValue,
                                function_ref<bool(const BaseStorage *)> isEqual,
-                               std::function<void(BaseStorage *)> cleanupFn) {
+                               function_ref<void(BaseStorage *)> cleanupFn) {
   impl->erase(kind, hashValue, isEqual, cleanupFn);
 }