Add a generic getValue to ElementsAttr for accessing a value at a given index.

PiperOrigin-RevId: 236013669

mlir/include/mlir/IR/Attributes.h

@@ -326,6 +326,10 @@ public:
 
   VectorOrTensorType getType() const;
 
+  /// Return the value at the given index. If index does not refer to a valid
+  /// element, then a null attribute is returned.
+  Attribute getValue(ArrayRef<uint64_t> index) const;
+
   /// Method for support type inquiry through isa, cast and dyn_cast.
   static bool kindof(Kind kind) {
     return kind >= Kind::FIRST_ELEMENTS_ATTR &&

mlir/lib/IR/Attributes.cpp

@@ -164,6 +164,24 @@ VectorOrTensorType ElementsAttr::getType() const {
   return static_cast<ImplType *>(attr)->type;
 }
 
+/// Return the value at the given index. If index does not refer to a valid
+/// element, then a null attribute is returned.
+Attribute ElementsAttr::getValue(ArrayRef<uint64_t> index) const {
+  switch (getKind()) {
+  case Attribute::Kind::SplatElements:
+    return cast<SplatElementsAttr>().getValue();
+  case Attribute::Kind::DenseFPElements:
+  case Attribute::Kind::DenseIntElements:
+    return cast<DenseElementsAttr>().getValue(index);
+  case Attribute::Kind::OpaqueElements:
+    return cast<OpaqueElementsAttr>().getValue(index);
+  case Attribute::Kind::SparseElements:
+    return cast<SparseElementsAttr>().getValue(index);
+  default:
+    llvm_unreachable("unknown ElementsAttr kind");
+  }
+}
+
 /// SplatElementsAttr
 
 Attribute SplatElementsAttr::getValue() const {

mlir/lib/StandardOps/StandardOps.cpp

@@ -1180,19 +1180,10 @@ Attribute ExtractElementOp::constantFold(ArrayRef<Attribute> operands,
     indices.push_back(indice.cast<IntegerAttr>().getInt());
   }
 
-  // Get the element value of the aggregate attribute with the given constant
-  // indices.
-  switch (aggregate.getKind()) {
-  case Attribute::Kind::DenseFPElements:
-  case Attribute::Kind::DenseIntElements:
-    return aggregate.cast<DenseElementsAttr>().getValue(indices);
-  case Attribute::Kind::OpaqueElements:
-    return aggregate.cast<OpaqueElementsAttr>().getValue(indices);
-  case Attribute::Kind::SparseElements:
-    return aggregate.cast<SparseElementsAttr>().getValue(indices);
-  default:
-    return Attribute();
-  }
+  // If this is an elements attribute, query the value at the given indices.
+  if (auto elementsAttr = aggregate.dyn_cast<ElementsAttr>())
+    return elementsAttr.getValue(indices);
+  return Attribute();
 }
 
 //===----------------------------------------------------------------------===//