Update ElementsAttr::isValidIndex to handle ElementsAttr with a scalar. Scalar will have rank 0.

[karimnosseir]

Summary

Update ElementsAttr::isValidIndex to handle ElementsAttr with a scalar. Scalar will have rank 0.

[rriddle]

Can you add a test for this?

[karimnosseir]

Added unit-test

[karimnosseir]

Thanks River. PTAL

[rriddle]

LGTM thanks.

mlir/unittests/IR/AttributeTest.cpp

195

[karimnosseir]

Updated test name

mlir/lib/IR/BuiltinAttributes.cpp

 }
 
 /// Return if the given 'index' refers to a valid element in this attribute.
 bool ElementsAttr::isValidIndex(ArrayRef<uint64_t> index) const {
   auto type = getType();
 
   // Verify that the rank of the indices matches the held type.
   auto 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.
   auto shape = type.getShape();
   return llvm::all_of(llvm::seq<int>(0, rank), [&](int i) {
     return static_cast<int64_t>(index[i]) < shape[i];
   });

mlir/unittests/IR/AttributeTest.cpp

 
 TEST(DenseComplexTest, ComplexAPIntSplat) {
   MLIRContext context;
   ComplexType complexType = ComplexType::get(IntegerType::get(&context, 64));
   std::complex<APInt> value(APInt(64, 10), APInt(64, 15));
   testSplat(complexType, value);
 }
 
+TEST(DenseScalarTest, ExtractZeroRankElement) {
+  MLIRContext context;
+  const int elementValue = 12;
+  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);
+}
+
 } // end namespace