[Mlir-commits] [mlir] 0d5caa8 - [mlir][DenseStringElementsAttr] Add support for the Attribute based get* methods.

[River Riddle]

Author: River Riddle
Date: 2020-05-01T16:34:35-07:00
New Revision: 0d5caa894017884fb02c9f3c1fd54abc0a9b0f06

URL: https://github.com/llvm/llvm-project/commit/0d5caa894017884fb02c9f3c1fd54abc0a9b0f06
DIFF: https://github.com/llvm/llvm-project/commit/0d5caa894017884fb02c9f3c1fd54abc0a9b0f06.diff

LOG: [mlir][DenseStringElementsAttr] Add support for the Attribute based get* methods.

This was missed in the original revision. This allows for using the opaque Attribute accessors when the elements are strings.

Added: 
    

Modified: 
    mlir/docs/LangRef.md
    mlir/include/mlir/IR/Attributes.h
    mlir/lib/IR/Attributes.cpp
    mlir/unittests/IR/AttributeTest.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/docs/LangRef.md b/mlir/docs/LangRef.md
index b3ee0c958c32..ff3af3e4e496 100644
--- a/mlir/docs/LangRef.md
+++ b/mlir/docs/LangRef.md
@@ -1334,9 +1334,10 @@ dense-elements-attribute ::= `dense` `<` attribute-value `>` `:`
 ```
 
 A dense elements attribute is an elements attribute where the storage for the
-constant vector or tensor value has been packed to the element bitwidth. The
-element type of the vector or tensor constant must be of integer, index, or
-floating point type.
+constant vector or tensor value has been densely packed. The attribute supports
+storing integer or floating point elements, with integer/index/floating element
+types. It also support storing string elements with a custom dialect string
+element type.
 
 ##### Opaque Elements Attribute
 

diff  --git a/mlir/include/mlir/IR/Attributes.h b/mlir/include/mlir/IR/Attributes.h
index 43628b8ce506..6d24cd087648 100644
--- a/mlir/include/mlir/IR/Attributes.h
+++ b/mlir/include/mlir/IR/Attributes.h
@@ -703,7 +703,8 @@ class DenseElementsAttr : public ElementsAttr {
 
   /// Constructs a dense elements attribute from an array of element values.
   /// Each element attribute value is expected to be an element of 'type'.
-  /// 'type' must be a vector or tensor with static shape.
+  /// 'type' must be a vector or tensor with static shape. If the element of
+  /// `type` is non-integer/index/float it is assumed to be a string type.
   static DenseElementsAttr get(ShapedType type, ArrayRef<Attribute> values);
 
   /// Constructs a dense integer elements attribute from an array of integer

diff  --git a/mlir/lib/IR/Attributes.cpp b/mlir/lib/IR/Attributes.cpp
index f51b273a9c89..f7540f6feffd 100644
--- a/mlir/lib/IR/Attributes.cpp
+++ b/mlir/lib/IR/Attributes.cpp
@@ -615,6 +615,8 @@ Attribute DenseElementsAttr::AttributeElementIterator::operator*() const {
     FloatElementIterator floatIt(floatEltTy.getFloatSemantics(), intIt);
     return FloatAttr::get(eltTy, *floatIt);
   }
+  if (owner.isa<DenseStringElementsAttr>())
+    return StringAttr::get(owner.getRawStringData()[index], eltTy);
   llvm_unreachable("unexpected element type");
 }
 
@@ -655,11 +657,23 @@ DenseElementsAttr::FloatElementIterator::FloatElementIterator(
 
 DenseElementsAttr DenseElementsAttr::get(ShapedType type,
                                          ArrayRef<Attribute> values) {
-  assert(type.getElementType().isIntOrIndexOrFloat() &&
-         "expected int or index or float element type");
   assert(hasSameElementsOrSplat(type, values));
 
+  // If the element type is not based on int/float/index, assume it is a string
+  // type.
   auto eltType = type.getElementType();
+  if (!type.getElementType().isIntOrIndexOrFloat()) {
+    SmallVector<StringRef, 8> stringValues;
+    stringValues.reserve(values.size());
+    for (Attribute attr : values) {
+      assert(attr.isa<StringAttr>() &&
+             "expected string value for non integer/index/float element");
+      stringValues.push_back(attr.cast<StringAttr>().getValue());
+    }
+    return get(type, stringValues);
+  }
+
+  // Otherwise, get the raw storage width to use for the allocation.
   size_t bitWidth = getDenseElementBitWidth(eltType);
   size_t storageBitWidth = getDenseElementStorageWidth(bitWidth);
 

diff  --git a/mlir/unittests/IR/AttributeTest.cpp b/mlir/unittests/IR/AttributeTest.cpp
index ad4b422eae91..7d22b5e5a07f 100644
--- a/mlir/unittests/IR/AttributeTest.cpp
+++ b/mlir/unittests/IR/AttributeTest.cpp
@@ -154,4 +154,12 @@ TEST(DenseSplatTest, StringSplat) {
   testSplat(stringType, value);
 }
 
+TEST(DenseSplatTest, StringAttrSplat) {
+  MLIRContext context;
+  Type stringType =
+      OpaqueType::get(Identifier::get("test", &context), "string", &context);
+  Attribute stringAttr = StringAttr::get("test-string", stringType);
+  testSplat(stringType, stringAttr);
+}
+
 } // end namespace