[mlir] Handle StringAttr in SparseElementsAttr::getZeroAttr.

[Chia-hungDuan]

Sorry, almost forgot this... Not sure if we need test for this, didn't see a proper place to do that.

[Mogball]

mlir/lib/IR/BuiltinAttributes.cpp

1340

auto? Since the type is already in the cast

[rriddle]

Can you add a test? An easy one would be to add a test for folding tensor.extract with a sparse string elements attribute (https://github.com/llvm/llvm-project/blob/b5e8348bf2ded134e45dab879a13d26f5b9c5815/mlir/test/Dialect/Tensor/canonicalize.mlir#L86)

LG otherwise.

[Chia-hungDuan]

Address review comment and fix string type. Will add test later.

[Chia-hungDuan]

Added test.

mlir/lib/IR/BuiltinAttributes.cpp

 /// Get a zero attribute for the given attribute type.
 Attribute SparseElementsAttr::getZeroAttr() const {
   auto eltType = getElementType();
 
   // Handle floating point elements.
   if (eltType.isa<FloatType>())
     return FloatAttr::get(eltType, 0);
 
+  // Handle string type.
+  if (getValues().isa<DenseStringElementsAttr>())
+    return StringAttr::get("", eltType);
+
   // Otherwise, this is an integer.
-  // TODO: Handle StringAttr here.
   return IntegerAttr::get(eltType, 0);
 }
 
 /// Flatten, and return, all of the sparse indices in this attribute in
 /// row-major order.
 std::vector<ptrdiff_t> SparseElementsAttr::getFlattenedSparseIndices() const {
   std::vector<ptrdiff_t> flatSparseIndices;
 

mlir/unittests/IR/AttributeTest.cpp

TEST(DenseScalarTest, ExtractZeroRankElement) {
   IntegerType intTy = IntegerType::get(&context, 32);
   Attribute value = IntegerAttr::get(intTy, elementValue);
   RankedTensorType shape = RankedTensorType::get({}, intTy);
 
   auto attr = DenseElementsAttr::get(shape, llvm::makeArrayRef({elementValue}));
   EXPECT_TRUE(attr.getValue({0}) == value);
 }
 
+TEST(SparseElementsAttrTest, GetZero) {
+  MLIRContext context;
+  context.allowUnregisteredDialects();
+
+  IntegerType intTy = IntegerType::get(&context, 32);
+  FloatType floatTy = FloatType::getF32(&context);
+  Type stringTy = OpaqueType::get(Identifier::get("test", &context), "string");
+
+  ShapedType tensorI32 = RankedTensorType::get({2, 2}, intTy);
+  ShapedType tensorF32 = RankedTensorType::get({2, 2}, floatTy);
+  ShapedType tensorString = RankedTensorType::get({2, 2}, stringTy);
+
+  auto indicesType =
+      RankedTensorType::get({1, 2}, IntegerType::get(&context, 64));
+  auto indices =
+      DenseIntElementsAttr::get(indicesType, {APInt(64, 0), APInt(64, 0)});
+
+  RankedTensorType intValueTy = RankedTensorType::get({1}, intTy);
+  auto intValue = DenseIntElementsAttr::get(intValueTy, {1});
+
+  RankedTensorType floatValueTy = RankedTensorType::get({1}, floatTy);
+  auto floatValue = DenseFPElementsAttr::get(floatValueTy, {1.0f});
+
+  RankedTensorType stringValueTy = RankedTensorType::get({1}, stringTy);
+  auto stringValue = DenseElementsAttr::get(stringValueTy, {StringRef("foo")});
+
+  auto sparseInt = SparseElementsAttr::get(tensorI32, indices, intValue);
+  auto sparseFloat = SparseElementsAttr::get(tensorF32, indices, floatValue);
+  auto sparseString =
+      SparseElementsAttr::get(tensorString, indices, stringValue);
+
+  // Only index (0, 0) contains an element, others are supposed to return
+  // the zero/empty value.
+  auto zeroIntValue = sparseInt.getValue({1, 1});
+  EXPECT_EQ(zeroIntValue.cast<IntegerAttr>().getInt(), 0);
+  EXPECT_TRUE(zeroIntValue.getType() == intTy);
+
+  auto zeroFloatValue = sparseFloat.getValue({1, 1});
+  EXPECT_EQ(zeroFloatValue.cast<FloatAttr>().getValueAsDouble(), 0.0f);
+  EXPECT_TRUE(zeroFloatValue.getType() == floatTy);
+
+  auto zeroStringValue = sparseString.getValue({1, 1});
+  EXPECT_TRUE(zeroStringValue.cast<StringAttr>().getValue().empty());
+  EXPECT_TRUE(zeroStringValue.getType() == stringTy);
+}
+
 } // end namespace