Add a new class NamedAttributeList to deduplicate named attribute handling between Function and Instruction.

PiperOrigin-RevId: 235830304

mlir/include/mlir/IR/Attributes.h

@@ -28,6 +28,7 @@ class Dialect;
 class Function;
 class FunctionAttr;
 class FunctionType;
+class Identifier;
 class IntegerSet;
 class Location;
 class MLIRContext;
@@ -54,6 +55,8 @@ struct DenseFPElementsAttributeStorage;
 struct OpaqueElementsAttributeStorage;
 struct SparseElementsAttributeStorage;
 
+class AttributeListStorage;
+
 } // namespace detail
 
 /// Attributes are known-constant values of operations and functions.
@@ -531,6 +534,46 @@ inline ::llvm::hash_code hash_value(Attribute arg) {
   return ::llvm::hash_value(arg.attr);
 }
 
+/// NamedAttribute is used for named attribute lists, it holds an identifier for
+/// the name and a value for the attribute. The attribute pointer should always
+/// be non-null.
+using NamedAttribute = std::pair<Identifier, Attribute>;
+
+/// A NamedAttributeList is used to manage a list of named attributes. This
+/// provides simple interfaces for adding/removing/finding attributes from
+/// within a raw AttributeListStorage.
+///
+/// We assume there will be relatively few attributes on a given function
+/// (maybe a dozen or so, but not hundreds or thousands) so we use linear
+/// searches for everything.
+class NamedAttributeList {
+public:
+  NamedAttributeList(MLIRContext *context, ArrayRef<NamedAttribute> attributes);
+
+  /// Return all of the attributes on this operation.
+  ArrayRef<NamedAttribute> getAttrs() const;
+
+  /// Replace the held attributes with ones provided in 'newAttrs'.
+  void setAttrs(MLIRContext *context, ArrayRef<NamedAttribute> attributes);
+
+  /// Return the specified attribute if present, null otherwise.
+  Attribute get(StringRef name) const;
+  Attribute get(Identifier name) const;
+
+  /// If the an attribute exists with the specified name, change it to the new
+  /// value.  Otherwise, add a new attribute with the specified name/value.
+  void set(MLIRContext *context, Identifier name, Attribute value);
+
+  enum class RemoveResult { Removed, NotFound };
+
+  /// Remove the attribute with the specified name if it exists.  The return
+  /// value indicates whether the attribute was present or not.
+  RemoveResult remove(MLIRContext *context, Identifier name);
+
+private:
+  detail::AttributeListStorage *attrs;
+};
+
 } // end namespace mlir.
 
 namespace llvm {

mlir/include/mlir/IR/Function.h

@@ -34,7 +34,6 @@
 #include "llvm/ADT/ilist.h"
 
 namespace mlir {
-class AttributeListStorage;
 class BlockAndValueMapping;
 class FunctionType;
 class MLIRContext;
@@ -42,11 +41,6 @@ class Module;
 template <typename ObjectType, typename ElementType> class ArgumentIterator;
 template <typename T> class OpPointer;
 
-/// NamedAttribute is used for function attribute lists, it holds an
-/// identifier for the name and a value for the attribute.  The attribute
-/// pointer should always be non-null.
-using NamedAttribute = std::pair<Identifier, Attribute>;
-
 /// This is the base class for all of the MLIR function types.
 class Function : public llvm::ilist_node_with_parent<Function, Module> {
 public:
@@ -66,9 +60,6 @@ public:
   /// Return the type of this function.
   FunctionType getType() const { return type; }
 
-  /// Returns all of the attributes on this function.
-  ArrayRef<NamedAttribute> getAttrs() const;
-
   MLIRContext *getContext() const;
   Module *getModule() { return module; }
   const Module *getModule() const { return module; }
@@ -179,6 +170,46 @@ public:
   const_args_iterator args_end() const;
   llvm::iterator_range<const_args_iterator> getArguments() const;
 
+  //===--------------------------------------------------------------------===//
+  // Attributes
+  //===--------------------------------------------------------------------===//
+
+  /// Functions may optionally carry a list of attributes that associate
+  /// constants to names.  Attributes may be dynamically added and removed over
+  /// the lifetime of an function.
+
+  /// Return all of the attributes on this instruction.
+  ArrayRef<NamedAttribute> getAttrs() const { return attrs.getAttrs(); }
+
+  /// Set the attributes held by this function.
+  void setAttrs(ArrayRef<NamedAttribute> attributes) {
+    attrs.setAttrs(getContext(), attributes);
+  }
+
+  /// Return the specified attribute if present, null otherwise.
+  Attribute getAttr(Identifier name) const { return attrs.get(name); }
+  Attribute getAttr(StringRef name) const { return attrs.get(name); }
+
+  template <typename AttrClass> AttrClass getAttrOfType(Identifier name) const {
+    return getAttr(name).dyn_cast_or_null<AttrClass>();
+  }
+
+  template <typename AttrClass> AttrClass getAttrOfType(StringRef name) const {
+    return getAttr(name).dyn_cast_or_null<AttrClass>();
+  }
+
+  /// If the an attribute exists with the specified name, change it to the new
+  /// value.  Otherwise, add a new attribute with the specified name/value.
+  void setAttr(Identifier name, Attribute value) {
+    attrs.set(getContext(), name, value);
+  }
+
+  /// Remove the attribute with the specified name if it exists.  The return
+  /// value indicates whether the attribute was present or not.
+  NamedAttributeList::RemoveResult removeAttr(Identifier name) {
+    return attrs.remove(getContext(), name);
+  }
+
   //===--------------------------------------------------------------------===//
   // Other
   //===--------------------------------------------------------------------===//
@@ -226,9 +257,6 @@ public:
   void cloneInto(Function *dest, BlockAndValueMapping &mapper) const;
 
 private:
-  /// Set the attributes held by this function.
-  void setAttributes(ArrayRef<NamedAttribute> attrs = {});
-
   /// The name of the function.
   Identifier name;
 
@@ -242,7 +270,7 @@ private:
   FunctionType type;
 
   /// This holds general named attributes for the function.
-  AttributeListStorage *attrs;
+  NamedAttributeList attrs;
 
   /// The contents of the body.
   BlockList blocks;

mlir/include/mlir/IR/Instruction.h

@@ -34,7 +34,6 @@
 #include "llvm/Support/TrailingObjects.h"
 
 namespace mlir {
-class AttributeListStorage;
 class BlockAndValueMapping;
 template <typename OpType> class ConstOpPointer;
 class Location;
@@ -389,31 +388,16 @@ public:
   // Attributes
   //===--------------------------------------------------------------------===//
 
-  // Operations may optionally carry a list of attributes that associate
+  // Instructions may optionally carry a list of attributes that associate
   // constants to names.  Attributes may be dynamically added and removed over
-  // the lifetime of an operation.
-  //
-  // We assume there will be relatively few attributes on a given operation
-  // (maybe a dozen or so, but not hundreds or thousands) so we use linear
-  // searches for everything.
+  // the lifetime of an instruction.
 
-  /// Return all of the attributes on this operation.
-  ArrayRef<NamedAttribute> getAttrs() const;
+  /// Return all of the attributes on this instruction.
+  ArrayRef<NamedAttribute> getAttrs() const { return attrs.getAttrs(); }
 
   /// Return the specified attribute if present, null otherwise.
-  Attribute getAttr(Identifier name) const {
-    for (auto elt : getAttrs())
-      if (elt.first == name)
-        return elt.second;
-    return nullptr;
-  }
-
-  Attribute getAttr(StringRef name) const {
-    for (auto elt : getAttrs())
-      if (elt.first.is(name))
-        return elt.second;
-    return nullptr;
-  }
+  Attribute getAttr(Identifier name) const { return attrs.get(name); }
+  Attribute getAttr(StringRef name) const { return attrs.get(name); }
 
   template <typename AttrClass> AttrClass getAttrOfType(Identifier name) const {
     return getAttr(name).dyn_cast_or_null<AttrClass>();
@@ -425,13 +409,15 @@ public:
 
   /// If the an attribute exists with the specified name, change it to the new
   /// value.  Otherwise, add a new attribute with the specified name/value.
-  void setAttr(Identifier name, Attribute value);
-
-  enum class RemoveResult { Removed, NotFound };
+  void setAttr(Identifier name, Attribute value) {
+    attrs.set(getContext(), name, value);
+  }
 
   /// Remove the attribute with the specified name if it exists.  The return
   /// value indicates whether the attribute was present or not.
-  RemoveResult removeAttr(Identifier name);
+  NamedAttributeList::RemoveResult removeAttr(Identifier name) {
+    return attrs.remove(getContext(), name);
+  }
 
   //===--------------------------------------------------------------------===//
   // Blocks
@@ -722,7 +708,7 @@ private:
   OperationName name;
 
   /// This holds general named attributes for the operation.
-  AttributeListStorage *attrs;
+  NamedAttributeList attrs;
 
   // allow ilist_traits access to 'block' field.
   friend struct llvm::ilist_traits<Instruction>;

mlir/include/mlir/IR/OperationSupport.h

@@ -166,11 +166,6 @@ private:
   const OperationProperties opProperties;
 };
 
-/// NamedAttribute is used for operation attribute lists, it holds an
-/// identifier for the name and a value for the attribute.  The attribute
-/// pointer should always be non-null.
-using NamedAttribute = std::pair<Identifier, Attribute>;
-
 class OperationName {
 public:
   using RepresentationUnion =

mlir/lib/IR/AttributeDetail.h

@@ -24,6 +24,7 @@
 
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/Attributes.h"
+#include "mlir/IR/Identifier.h"
 #include "mlir/IR/IntegerSet.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/IR/StandardTypes.h"
@@ -148,6 +149,36 @@ struct SparseElementsAttributeStorage : public ElementsAttributeStorage {
   DenseElementsAttr values;
 };
 
+/// A raw list of named attributes stored as a trailing array.
+class AttributeListStorage final
+    : private llvm::TrailingObjects<AttributeListStorage, NamedAttribute> {
+  friend class llvm::TrailingObjects<AttributeListStorage, NamedAttribute>;
+
+public:
+  /// Given a list of NamedAttribute's, canonicalize the list (sorting
+  /// by name) and return the unique'd result.  Note that the empty list is
+  /// represented with a null pointer.
+  static AttributeListStorage *get(ArrayRef<NamedAttribute> attrs,
+                                   MLIRContext *context);
+
+  /// Return the element constants for this aggregate constant.  These are
+  /// known to all be constants.
+  ArrayRef<NamedAttribute> getElements() const {
+    return {getTrailingObjects<NamedAttribute>(), numElements};
+  }
+
+private:
+  // This is used by the llvm::TrailingObjects base class.
+  size_t numTrailingObjects(OverloadToken<NamedAttribute>) const {
+    return numElements;
+  }
+  AttributeListStorage() = delete;
+  AttributeListStorage(const AttributeListStorage &) = delete;
+  AttributeListStorage(unsigned numElements) : numElements(numElements) {}
+
+  /// This is the number of attributes.
+  const unsigned numElements;
+};
 } // namespace detail
 } // namespace mlir
 

mlir/lib/IR/AttributeListStorage.h

@@ -1,58 +0,0 @@
-//===- AttributeListStorage.h - Attr representation for ops -----*- C++ -*-===//
-//
-// Copyright 2019 The MLIR Authors.
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//   http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-// =============================================================================
-
-#ifndef ATTRIBUTELISTSTORAGE_H
-#define ATTRIBUTELISTSTORAGE_H
-
-#include "mlir/IR/OperationSupport.h"
-#include "llvm/Support/TrailingObjects.h"
-
-namespace mlir {
-
-class AttributeListStorage final
-    : private llvm::TrailingObjects<AttributeListStorage, NamedAttribute> {
-  friend class llvm::TrailingObjects<AttributeListStorage, NamedAttribute>;
-
-public:
-  /// Given a list of NamedAttribute's, canonicalize the list (sorting
-  /// by name) and return the unique'd result.  Note that the empty list is
-  /// represented with a null pointer.
-  static AttributeListStorage *get(ArrayRef<NamedAttribute> attrs,
-                                   MLIRContext *context);
-
-  /// Return the element constants for this aggregate constant.  These are
-  /// known to all be constants.
-  ArrayRef<NamedAttribute> getElements() const {
-    return {getTrailingObjects<NamedAttribute>(), numElements};
-  }
-
-private:
-  // This is used by the llvm::TrailingObjects base class.
-  size_t numTrailingObjects(OverloadToken<NamedAttribute>) const {
-    return numElements;
-  }
-  AttributeListStorage() = delete;
-  AttributeListStorage(const AttributeListStorage &) = delete;
-  AttributeListStorage(unsigned numElements) : numElements(numElements) {}
-
-  /// This is the number of attributes.
-  const unsigned numElements;
-};
-
-} // end namespace mlir
-
-#endif

mlir/lib/IR/Attributes.cpp

@@ -447,3 +447,80 @@ Attribute SparseElementsAttr::getValue(ArrayRef<uint64_t> index) const {
   // Otherwise, return the held sparse value element.
   return getValues().getValue(it->second);
 }
+
+/// NamedAttributeList
+
+NamedAttributeList::NamedAttributeList(MLIRContext *context,
+                                       ArrayRef<NamedAttribute> attributes) {
+  setAttrs(context, attributes);
+}
+
+/// Return all of the attributes on this operation.
+ArrayRef<NamedAttribute> NamedAttributeList::getAttrs() const {
+  return attrs ? attrs->getElements() : llvm::None;
+}
+
+/// Replace the held attributes with ones provided in 'newAttrs'.
+void NamedAttributeList::setAttrs(MLIRContext *context,
+                                  ArrayRef<NamedAttribute> attributes) {
+  // Don't create an attribute list if there are no attributes.
+  if (attributes.empty()) {
+    attrs = nullptr;
+    return;
+  }
+
+  assert(llvm::all_of(attributes,
+                      [](const NamedAttribute &attr) { return attr.second; }) &&
+         "attributes cannot have null entries");
+  attrs = AttributeListStorage::get(attributes, context);
+}
+
+/// Return the specified attribute if present, null otherwise.
+Attribute NamedAttributeList::get(StringRef name) const {
+  for (auto elt : getAttrs())
+    if (elt.first.is(name))
+      return elt.second;
+  return nullptr;
+}
+Attribute NamedAttributeList::get(Identifier name) const {
+  return get(name.strref());
+}
+
+/// If the an attribute exists with the specified name, change it to the new
+/// value.  Otherwise, add a new attribute with the specified name/value.
+void NamedAttributeList::set(MLIRContext *context, Identifier name,
+                             Attribute value) {
+  assert(value && "attributes may never be null");
+
+  // If we already have this attribute, replace it.
+  auto origAttrs = getAttrs();
+  SmallVector<NamedAttribute, 8> newAttrs(origAttrs.begin(), origAttrs.end());
+  for (auto &elt : newAttrs)
+    if (elt.first == name) {
+      elt.second = value;
+      attrs = AttributeListStorage::get(newAttrs, context);
+      return;
+    }
+
+  // Otherwise, add it.
+  newAttrs.push_back({name, value});
+  attrs = AttributeListStorage::get(newAttrs, context);
+}
+
+/// Remove the attribute with the specified name if it exists.  The return
+/// value indicates whether the attribute was present or not.
+auto NamedAttributeList::remove(MLIRContext *context, Identifier name)
+    -> RemoveResult {
+  auto origAttrs = getAttrs();
+  for (unsigned i = 0, e = origAttrs.size(); i != e; ++i) {
+    if (origAttrs[i].first == name) {
+      SmallVector<NamedAttribute, 8> newAttrs;
+      newAttrs.reserve(origAttrs.size() - 1);
+      newAttrs.append(origAttrs.begin(), origAttrs.begin() + i);
+      newAttrs.append(origAttrs.begin() + i + 1, origAttrs.end());
+      attrs = AttributeListStorage::get(newAttrs, context);
+      return RemoveResult::Removed;
+    }
+  }
+  return RemoveResult::NotFound;
+}

mlir/lib/IR/Function.cpp

@@ -16,7 +16,6 @@
 // =============================================================================
 
 #include "mlir/IR/Function.h"
-#include "AttributeListStorage.h"
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/BlockAndValueMapping.h"
 #include "mlir/IR/MLIRContext.h"
@@ -30,9 +29,7 @@ using namespace mlir;
 Function::Function(Location location, StringRef name, FunctionType type,
                    ArrayRef<NamedAttribute> attrs)
     : name(Identifier::get(name, type.getContext())), location(location),
-      type(type), blocks(this) {
-  setAttributes(attrs);
-}
+      type(type), attrs(type.getContext(), attrs), blocks(this) {}
 
 Function::~Function() {
   // Instructions may have cyclic references, which need to be dropped before we
@@ -44,13 +41,6 @@ Function::~Function() {
   FunctionAttr::dropFunctionReference(this);
 }
 
-ArrayRef<NamedAttribute> Function::getAttrs() const {
-  if (attrs)
-    return attrs->getElements();
-  else
-    return {};
-}
-
 MLIRContext *Function::getContext() const { return getType().getContext(); }
 
 Module *llvm::ilist_traits<Function>::getContainingModule() {
@@ -160,7 +150,7 @@ void Function::cloneInto(Function *dest, BlockAndValueMapping &mapper) const {
     assert((insertPair.second || insertPair.first->second == attr.second) &&
            "the two functions have incompatible attributes");
   }
-  dest->setAttributes(newAttrs.takeVector());
+  dest->setAttrs(newAttrs.takeVector());
 
   // Clone the block list.
   blocks.cloneInto(&dest->blocks, mapper, dest->getContext());
@@ -195,11 +185,6 @@ Function *Function::clone() const {
   return clone(mapper);
 }
 
-/// Set the attributes held by this function.
-void Function::setAttributes(ArrayRef<NamedAttribute> attrs) {
-  this->attrs = AttributeListStorage::get(attrs, getContext());
-}
-
 //===----------------------------------------------------------------------===//
 // Function implementation.
 //===----------------------------------------------------------------------===//

mlir/lib/IR/Instruction.cpp

@@ -16,7 +16,6 @@
 // =============================================================================
 
 #include "mlir/IR/Instruction.h"
-#include "AttributeListStorage.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/BlockAndValueMapping.h"
@@ -243,12 +242,7 @@ Instruction::Instruction(Location location, OperationName name,
                          ArrayRef<NamedAttribute> attributes,
                          MLIRContext *context)
     : location(location), numResults(numResults), numSuccs(numSuccessors),
-      numBlockLists(numBlockLists), name(name) {
-  assert(llvm::all_of(attributes,
-                      [](const NamedAttribute &attr) { return attr.second; }) &&
-         "Attributes cannot have null entries");
-  this->attrs = AttributeListStorage::get(attributes, context);
-}
+      numBlockLists(numBlockLists), name(name), attrs(context, attributes) {}
 
 // Instructions are deleted through the destroy() member because they are
 // allocated via malloc.
@@ -487,12 +481,6 @@ bool Instruction::use_empty() const {
   return true;
 }
 
-ArrayRef<NamedAttribute> Instruction::getAttrs() const {
-  if (!attrs)
-    return {};
-  return attrs->getElements();
-}
-
 bool Instruction::isReturn() const { return isa<ReturnOp>(); }
 
 void Instruction::setSuccessor(Block *block, unsigned index) {
@@ -523,45 +511,6 @@ auto Instruction::getSuccessorOperands(unsigned index) -> operand_range {
                            succOperandIndex + getNumSuccessorOperands(index))};
 }
 
-/// If an attribute exists with the specified name, change it to the new
-/// value.  Otherwise, add a new attribute with the specified name/value.
-void Instruction::setAttr(Identifier name, Attribute value) {
-  assert(value && "attributes may never be null");
-  auto origAttrs = getAttrs();
-
-  SmallVector<NamedAttribute, 8> newAttrs(origAttrs.begin(), origAttrs.end());
-  auto *context = getContext();
-
-  // If we already have this attribute, replace it.
-  for (auto &elt : newAttrs)
-    if (elt.first == name) {
-      elt.second = value;
-      attrs = AttributeListStorage::get(newAttrs, context);
-      return;
-    }
-
-  // Otherwise, add it.
-  newAttrs.push_back({name, value});
-  attrs = AttributeListStorage::get(newAttrs, context);
-}
-
-/// Remove the attribute with the specified name if it exists.  The return
-/// value indicates whether the attribute was present or not.
-auto Instruction::removeAttr(Identifier name) -> RemoveResult {
-  auto origAttrs = getAttrs();
-  for (unsigned i = 0, e = origAttrs.size(); i != e; ++i) {
-    if (origAttrs[i].first == name) {
-      SmallVector<NamedAttribute, 8> newAttrs;
-      newAttrs.reserve(origAttrs.size() - 1);
-      newAttrs.append(origAttrs.begin(), origAttrs.begin() + i);
-      newAttrs.append(origAttrs.begin() + i + 1, origAttrs.end());
-      attrs = AttributeListStorage::get(newAttrs, getContext());
-      return RemoveResult::Removed;
-    }
-  }
-  return RemoveResult::NotFound;
-}
-
 /// Attempt to constant fold this operation with the specified constant
 /// operand values.  If successful, this returns false and fills in the
 /// results vector.  If not, this returns true and results is unspecified.

mlir/lib/IR/MLIRContext.cpp

@@ -19,7 +19,6 @@
 #include "AffineExprDetail.h"
 #include "AffineMapDetail.h"
 #include "AttributeDetail.h"
-#include "AttributeListStorage.h"
 #include "IntegerSetDetail.h"
 #include "LocationDetail.h"
 #include "TypeDetail.h"