Add constant folding for ExtractElementOp when the aggregate is an OpaqueElementsAttr.

PiperOrigin-RevId: 235533283

mlir/include/mlir/IR/Attributes.h

@@ -438,6 +438,10 @@ public:
 
   StringRef getValue() 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;
+
   /// Decodes the attribute value using dialect-specific decoding hook.
   /// Returns false if decoding is successful. If not, returns true and leaves
   /// 'result' argument unspecified.

mlir/lib/IR/Attributes.cpp

@@ -342,6 +342,14 @@ StringRef OpaqueElementsAttr::getValue() const {
   return static_cast<ImplType *>(attr)->bytes;
 }
 
+/// Return the value at the given index. If index does not refer to a valid
+/// element, then a null attribute is returned.
+Attribute OpaqueElementsAttr::getValue(ArrayRef<uint64_t> index) const {
+  if (Dialect *dialect = getDialect())
+    return dialect->extractElementHook(*this, index);
+  return Attribute();
+}
+
 Dialect *OpaqueElementsAttr::getDialect() const {
   return static_cast<ImplType *>(attr)->dialect;
 }

mlir/lib/StandardOps/StandardOps.cpp

@@ -1186,6 +1186,8 @@ Attribute ExtractElementOp::constantFold(ArrayRef<Attribute> operands,
   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: