[mlir] Make ElementsAttr inherit from TypedAttr

[rkayaith]

Summary

This allows implicit conversion from ElementsAttr to TypedAttr, but
required renaming the ElementsAttr::getType() interface method to
getShapedType.

[rkayaith]

mlir/include/mlir/IR/BuiltinAttributeInterfaces.td

182

I'm not sure this is the best name, open to suggestions.

mlir/include/mlir/Dialect/CommonFolders.h

if (operands[0].isa<SplatElementsAttr>()) {
     SmallVector<ElementValueT> elementResults;
     elementResults.reserve(op.getNumElements());
     for (size_t i = 0, e = op.getNumElements(); i < e; ++i, ++opIt) {
       auto elementResult = calculate(*opIt);
       if (!elementResult)
         return {};
       elementResults.push_back(*elementResult);
     }
-    return DenseElementsAttr::get(op.getType(), elementResults);
+    return DenseElementsAttr::get(op.getShapedType(), elementResults);
   }
   return {};
 }
 
 template <class AttrElementT,
           class ElementValueT = typename AttrElementT::ValueType,
           class CalculationT = function_ref<ElementValueT(ElementValueT)>>
 Attribute constFoldUnaryOp(ArrayRef<Attribute> operands,

mlir/include/mlir/IR/BuiltinAttributeInterfaces.td

 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_IR_BUILTINATTRIBUTEINTERFACES_TD_
 #define MLIR_IR_BUILTINATTRIBUTEINTERFACES_TD_
 
 include "mlir/IR/OpBase.td"
 
 //===----------------------------------------------------------------------===//
+// TypedAttrInterface
+//===----------------------------------------------------------------------===//
+
+def TypedAttrInterface : AttrInterface<"TypedAttr"> {
+  let cppNamespace = "::mlir";
+
+  let description = [{
+    This interface is used for attributes that have a type. The type of an
+    attribute is understood to represent the type of the data contained in the
+    attribute and is often used as the type of a value with this data.
+  }];
+
+  let methods = [InterfaceMethod<
+    "Get the attribute's type",
+    "::mlir::Type", "getType"
+  >];
+}
+
+//===----------------------------------------------------------------------===//
 // ElementsAttrInterface
 //===----------------------------------------------------------------------===//
 
-def ElementsAttrInterface : AttrInterface<"ElementsAttr"> {
+def ElementsAttrInterface : AttrInterface<"ElementsAttr", [TypedAttrInterface]> {
   let cppNamespace = "::mlir";
   let description = [{
     This interface is used for attributes that contain the constant elements of
     a tensor or vector type. It allows for opaquely interacting with the
     elements of the underlying attribute, and most importantly allows for
     accessing the element values (including iteration) in any of the C++ data
     types supported by the underlying attribute.
 
let description = [{
     auto value_begin_impl(OverloadToken<llvm::APInt>) const {
       auto it = llvm::map_range(getElements(), [=](uint64_t value) {
         return llvm::APInt(/*numBits=*/64, value);
       }).begin();
       return FailureOr<decltype(it)>(std::move(it));
     }
     /// * Attribute
     auto value_begin_impl(OverloadToken<mlir::Attribute>) const {
-      mlir::Type elementType = getType().getElementType();
+      mlir::Type elementType = getShapedType().getElementType();
       auto it = llvm::map_range(getElements(), [=](uint64_t value) {
         return mlir::IntegerAttr::get(elementType,
                                       llvm::APInt(/*numBits=*/64, value));
       }).begin();
       return FailureOr<decltype(it)>(std::move(it));
     }
     ```
 
(ins "::mlir::TypeID":$elementID), [{}], /*defaultImplementation=*/[{
 
       return getValueImpl(
         (typename ConcreteAttr::NonContiguousIterableTypesT *)nullptr,
         elementID, /*isContiguous=*/std::false_type());
     }]>,
     InterfaceMethod<[{
       Returns true if the attribute elements correspond to a splat, i.e. that
       all elements of the attribute are the same value.
-    }], "bool", "isSplat", (ins), /*defaultImplementation=*/[{}], [{
+    }], "bool", "isSplat", (ins), [{}], /*defaultImplementation=*/[{
         // By default, only check for a single element splat.
         return $_attr.getNumElements() == 1;
     }]>,
     InterfaceMethod<[{
       Returns the shaped type of the elements attribute.
-    }], "::mlir::ShapedType", "getType">
+    }], "::mlir::ShapedType", "getShapedType", (ins), [{}], /*defaultImplementation=*/[{
+        return $_attr.getType();
+    }]>
   ];
 
   string ElementsAttrInterfaceAccessors = [{
     /// Return the number of elements held by this attribute.
     int64_t size() const { return getNumElements(); }
 
     /// Return if the attribute holds no elements.
     bool empty() const { return size() == 0; }
let extraClassDeclaration = [{
     /// multi-dimensional index.
     uint64_t getFlattenedIndex(ArrayRef<uint64_t> index) const {
       return getFlattenedIndex(*this, index);
     }
     static uint64_t getFlattenedIndex(Type type,
                                       ArrayRef<uint64_t> index);
     static uint64_t getFlattenedIndex(ElementsAttr elementsAttr,
                                       ArrayRef<uint64_t> index) {
-      return getFlattenedIndex(elementsAttr.getType(), index);
+      return getFlattenedIndex(elementsAttr.getShapedType(), index);
     }
 
     /// Returns the number of elements held by this attribute.
     int64_t getNumElements() const { return getNumElements(*this); }
     static int64_t getNumElements(ElementsAttr elementsAttr);
 
     //===------------------------------------------------------------------===//
     // Value Iteration
let extraClassDeclaration = [{
     T getSplatValue() const {
       assert(isSplat() && "expected splat attribute");
       return *value_begin<T>();
     }
 
     /// Return the elements of this attribute as a value of type 'T'.
     template <typename T>
     DefaultValueCheckT<T, iterator_range<T>> getValues() const {
-      return {getType(), value_begin<T>(), value_end<T>()};
+      return {getShapedType(), value_begin<T>(), value_end<T>()};
     }
     template <typename T>
     DefaultValueCheckT<T, iterator<T>> value_begin() const;
     template <typename T>
     DefaultValueCheckT<T, iterator<T>> value_end() const {
       return iterator<T>({}, size());
     }
 
     /// Return the held element values a range of T, where T is a derived
     /// attribute type.
     template <typename T>
     using DerivedAttrValueIterator =
       llvm::mapped_iterator<iterator<Attribute>, T (*)(Attribute)>;
     template <typename T>
     using DerivedAttrValueIteratorRange =
       detail::ElementsAttrRange<DerivedAttrValueIterator<T>>;
     template <typename T, typename = DerivedAttrValueCheckT<T>>
     DerivedAttrValueIteratorRange<T> getValues() const {
       auto castFn = [](Attribute attr) { return attr.template cast<T>(); };
-      return {getType(), llvm::map_range(getValues<Attribute>(),
+      return {getShapedType(), llvm::map_range(getValues<Attribute>(),
               static_cast<T (*)(Attribute)>(castFn))};
     }
     template <typename T, typename = DerivedAttrValueCheckT<T>>
     DerivedAttrValueIterator<T> value_begin() const {
       return getValues<T>().begin();
     }
     template <typename T, typename = DerivedAttrValueCheckT<T>>
     DerivedAttrValueIterator<T> value_end() const {
       return {value_end<Attribute>(), nullptr};
     }
 
     //===------------------------------------------------------------------===//
     // Failable Value Iteration
 
     /// If this attribute supports iterating over element values of type `T`,
     /// return the iterable range. Otherwise, return std::nullopt.
     template <typename T>
     DefaultValueCheckT<T, std::optional<iterator_range<T>>> tryGetValues() const {
       if (std::optional<iterator<T>> beginIt = try_value_begin<T>())
-        return iterator_range<T>(getType(), *beginIt, value_end<T>());
+        return iterator_range<T>(getShapedType(), *beginIt, value_end<T>());
       return std::nullopt;
     }
     template <typename T>
     DefaultValueCheckT<T, std::optional<iterator<T>>> try_value_begin() const;
 
     /// If this attribute supports iterating over element values of type `T`,
     /// return the iterable range. Otherwise, return std::nullopt.
     template <typename T, typename = DerivedAttrValueCheckT<T>>
     std::optional<DerivedAttrValueIteratorRange<T>> tryGetValues() const {
       auto values = tryGetValues<Attribute>();
       if (!values)
         return std::nullopt;
 
       auto castFn = [](Attribute attr) { return attr.template cast<T>(); };
       return DerivedAttrValueIteratorRange<T>(
-        getType(),
+        getShapedType(),
         llvm::map_range(*values, static_cast<T (*)(Attribute)>(castFn))
       );
     }
     template <typename T, typename = DerivedAttrValueCheckT<T>>
     std::optional<DerivedAttrValueIterator<T>> try_value_begin() const {
       if (auto values = tryGetValues<T>())
         return values->begin();
       return std::nullopt;
InterfaceMethod<
       [{
         return ::mlir::detail::verifyAffineMapAsLayout($_attr.getAffineMap(),
                                                        shape, emitError);
       }]
     >
   ];
 }
 
-//===----------------------------------------------------------------------===//
-// TypedAttrInterface
-//===----------------------------------------------------------------------===//
-
-def TypedAttrInterface : AttrInterface<"TypedAttr"> {
-  let cppNamespace = "::mlir";
-
-  let description = [{
-    This interface is used for attributes that have a type. The type of an
-    attribute is understood to represent the type of the data contained in the
-    attribute and is often used as the type of a value with this data.
-  }];
-
-  let methods = [InterfaceMethod<
-    "Get the attribute's type",
-    "::mlir::Type", "getType"
-  >];
-}
-
 #endif // MLIR_IR_BUILTINATTRIBUTEINTERFACES_TD_

mlir/include/mlir/IR/BuiltinAttributes.td

def Builtin_DenseArray : Builtin_Attr<"DenseArray"> {
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // DenseIntOrFPElementsAttr
 //===----------------------------------------------------------------------===//
 
 def Builtin_DenseIntOrFPElementsAttr : Builtin_Attr<
-    "DenseIntOrFPElements", [ElementsAttrInterface, TypedAttrInterface],
+    "DenseIntOrFPElements", [ElementsAttrInterface],
     "DenseElementsAttr"
   > {
   let summary = "An Attribute containing a dense multi-dimensional array of "
                 "integer or floating-point values";
   let description = [{
     Syntax:
 
     ```
def Builtin_DenseIntOrFPElementsAttr : Builtin_Attr<
   let skipDefaultBuilders = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // DenseStringElementsAttr
 //===----------------------------------------------------------------------===//
 
 def Builtin_DenseStringElementsAttr : Builtin_Attr<
-    "DenseStringElements", [ElementsAttrInterface, TypedAttrInterface],
+    "DenseStringElements", [ElementsAttrInterface],
     "DenseElementsAttr"
   > {
   let summary = "An Attribute containing a dense multi-dimensional array of "
                 "strings";
   let description = [{
     Syntax:
 
     ```
def Builtin_DenseStringElementsAttr : Builtin_Attr<
   let skipDefaultBuilders = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // DenseResourceElementsAttr
 //===----------------------------------------------------------------------===//
 
 def Builtin_DenseResourceElementsAttr : Builtin_Attr<"DenseResourceElements", [
-    ElementsAttrInterface, TypedAttrInterface
+    ElementsAttrInterface
   ]> {
   let summary = "An Attribute containing a dense multi-dimensional array "
                 "backed by a resource";
   let description = [{
     Syntax:
 
     ```
     dense-resource-elements-attribute ::=
def Builtin_OpaqueAttr : Builtin_Attr<"Opaque", [TypedAttrInterface]> {
   let skipDefaultBuilders = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // SparseElementsAttr
 //===----------------------------------------------------------------------===//
 
 def Builtin_SparseElementsAttr : Builtin_Attr<
-    "SparseElements", [ElementsAttrInterface, TypedAttrInterface]
+    "SparseElements", [ElementsAttrInterface]
   > {
   let summary = "An opaque representation of a multi-dimensional array";
   let description = [{
     Syntax:
 
     ```
     sparse-elements-attribute ::= `sparse` `<` attribute-value `,`
                                   attribute-value `>` `:`

mlir/lib/IR/BuiltinAttributeInterfaces.cpp

 
 #include "mlir/IR/BuiltinAttributeInterfaces.cpp.inc"
 
 //===----------------------------------------------------------------------===//
 // ElementsAttr
 //===----------------------------------------------------------------------===//
 
 Type ElementsAttr::getElementType(ElementsAttr elementsAttr) {
-  return elementsAttr.getType().getElementType();
+  return elementsAttr.getShapedType().getElementType();
 }
 
 int64_t ElementsAttr::getNumElements(ElementsAttr elementsAttr) {
-  return elementsAttr.getType().getNumElements();
+  return elementsAttr.getShapedType().getNumElements();
 }
 
 bool ElementsAttr::isValidIndex(ShapedType type, ArrayRef<uint64_t> index) {
   // Verify that the rank of the indices matches the held type.
   int64_t rank = type.getRank();
   if (rank == 0 && index.size() == 1 && index[0] == 0)
     return true;
   if (rank != static_cast<int64_t>(index.size()))
     return false;
 
   // Verify that all of the indices are within the shape dimensions.
   ArrayRef<int64_t> shape = type.getShape();
   return llvm::all_of(llvm::seq<int>(0, rank), [&](int i) {
     int64_t dim = static_cast<int64_t>(index[i]);
     return 0 <= dim && dim < shape[i];
   });
 }
 bool ElementsAttr::isValidIndex(ElementsAttr elementsAttr,
                                 ArrayRef<uint64_t> index) {
-  return isValidIndex(elementsAttr.getType(), index);
+  return isValidIndex(elementsAttr.getShapedType(), index);
 }
 
 uint64_t ElementsAttr::getFlattenedIndex(Type type, ArrayRef<uint64_t> index) {
   ShapedType shapeType = type.cast<ShapedType>();
   assert(isValidIndex(shapeType, index) &&
          "expected valid multi-dimensional index");
 
   // Reduce the provided multidimensional index into a flattended 1D row-major

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

llvm::Constant *mlir::LLVM::detail::getLLVMConstant(
   if (llvm::Constant *result =
           convertDenseElementsAttr(loc, attr.dyn_cast<DenseElementsAttr>(),
                                    llvmType, moduleTranslation)) {
     return result;
   }
 
   // Fall back to element-by-element construction otherwise.
   if (auto elementsAttr = attr.dyn_cast<ElementsAttr>()) {
-    assert(elementsAttr.getType().hasStaticShape());
-    assert(!elementsAttr.getType().getShape().empty() &&
+    assert(elementsAttr.getShapedType().hasStaticShape());
+    assert(!elementsAttr.getShapedType().getShape().empty() &&
            "unexpected empty elements attribute shape");
 
     SmallVector<llvm::Constant *, 8> constants;
     constants.reserve(elementsAttr.getNumElements());
     llvm::Type *innermostType = getInnermostElementType(llvmType);
     for (auto n : elementsAttr.getValues<Attribute>()) {
       constants.push_back(
           getLLVMConstant(innermostType, n, loc, moduleTranslation));
       if (!constants.back())
         return nullptr;
     }
     ArrayRef<llvm::Constant *> constantsRef = constants;
     llvm::Constant *result = buildSequentialConstant(
-        constantsRef, elementsAttr.getType().getShape(), llvmType, loc);
+        constantsRef, elementsAttr.getShapedType().getShape(), llvmType, loc);
     assert(constantsRef.empty() && "did not consume all elemental constants");
     return result;
   }
 
   if (auto stringAttr = attr.dyn_cast<StringAttr>()) {
     return llvm::ConstantDataArray::get(
         moduleTranslation.getLLVMContext(),
         ArrayRef<char>{stringAttr.getValue().data(),

mlir/lib/Transforms/ViewOpGraph.cpp

void emitMlirAttr(raw_ostream &os, Attribute attr) {
     if (attr.isa<SplatElementsAttr>()) {
       attr.print(os);
       return;
     }
 
     // Elide "big" elements attributes.
     auto elements = attr.dyn_cast<ElementsAttr>();
     if (elements && elements.getNumElements() > largeAttrLimit) {
-      os << std::string(elements.getType().getRank(), '[') << "..."
-         << std::string(elements.getType().getRank(), ']') << " : "
+      os << std::string(elements.getShapedType().getRank(), '[') << "..."
+         << std::string(elements.getShapedType().getRank(), ']') << " : "
          << elements.getType();
       return;
     }
 
     auto array = attr.dyn_cast<ArrayAttr>();
     if (array && static_cast<int64_t>(array.size()) > largeAttrLimit) {
       os << "[...]";
       return;

mlir/test/lib/Dialect/Test/TestAttrDefs.td

 def TestAttrTrait : NativeAttrTrait<"TestAttrTrait">;
 
 // The definition of a singleton attribute that has a trait.
 def AttrWithTrait : Test_Attr<"AttrWithTrait", [TestAttrTrait]> {
   let mnemonic = "attr_with_trait";
 }
 
 // Test support for ElementsAttrInterface.
-def TestI64ElementsAttr : Test_Attr<"TestI64Elements", [
-    ElementsAttrInterface, TypedAttrInterface
-  ]> {
+def TestI64ElementsAttr : Test_Attr<"TestI64Elements", [ElementsAttrInterface]> {
   let mnemonic = "i64_elements";
   let parameters = (ins
     AttributeSelfTypeParameter<"", "::mlir::ShapedType">:$type,
     ArrayRefParameter<"uint64_t">:$elements
   );
   let extraClassDeclaration = [{
     /// The set of data types that can be iterated by this attribute.
     using ContiguousIterableTypesT = std::tuple<uint64_t>;
def TestOverrideBuilderAttr : Test_Attr<"TestOverrideBuilder"> {
   let skipDefaultBuilders = 1;
   let genVerifyDecl = 1;
   let builders = [AttrBuilder<(ins "int":$a), [{
     return ::mlir::IntegerAttr::get(::mlir::IndexType::get($_ctxt), a);
   }], "::mlir::Attribute">];
 }
 
 // Test simple extern 1D vector using ElementsAttrInterface.
-def TestExtern1DI64ElementsAttr : Test_Attr<"TestExtern1DI64Elements", [
-    ElementsAttrInterface, TypedAttrInterface
-  ]> {
+def TestExtern1DI64ElementsAttr : Test_Attr<"TestExtern1DI64Elements", [ElementsAttrInterface]> {
   let mnemonic = "e1di64_elements";
   let parameters = (ins
     AttributeSelfTypeParameter<"", "::mlir::ShapedType">:$type,
     ResourceHandleParameter<"TestDialectResourceBlobHandle">:$handle
   );
   let extraClassDeclaration = [{
     /// Return the elements referenced by this attribute.
     llvm::ArrayRef<uint64_t> getElements() const;