[mlir] Fix a crash of attribute tests when given a non-integer type

[sotoshigoto]

Summary

The crash issue occurs in the TestBuiltinAttributeInterface when a
dense attribute with non-integer elements is provided. The tryGetValues
function should first check conditions, including casting T from the
element's type and verifying whether the elements are empty or not. This
fix resolves the crash issue reported in:

• https://github.com/llvm/llvm-project/issues/60581
• https://github.com/llvm/llvm-project/issues/61871

[rkayaith]

mlir/test/IR/elements-attr-interface.mlir

33–34

I'm not too familiar with this interface, but this change doesn't seem right to me. Successfully returning an empty range seems like the desired behaviour.

[sotoshigoto]

mlir/test/IR/elements-attr-interface.mlir

33–34

Thanks for your review! I'll address the empty range behavior by checking for attribute casting on empty elements.

[sotoshigoto]

addressed review comment

[sotoshigoto]

retry to upload patch

mlir/include/mlir/IR/BuiltinAttributeInterfaces.td

let extraClassDeclaration = [{
     /// If this attribute supports iterating over element values of type `T`,
     /// return the iterable range. Otherwise, return std::nullopt.
     template <typename T, typename = DerivedAttrValueCheckT<T>>
     std::optional<DerivedAttrValueIteratorRange<T>> tryGetValues() const {
       auto values = tryGetValues<Attribute>();
       if (!values)
         return std::nullopt;
 
+      // Check that an element can be casted as 'T', which is a derived attribute.
+      if (!empty())
+        if (!(*(value_begin<Attribute>())).dyn_cast<T>())
+          return std::nullopt;
+
       auto castFn = [](Attribute attr) { return ::llvm::cast<T>(attr); };
       return DerivedAttrValueIteratorRange<T>(
         getShapedType(),
         llvm::map_range(*values, static_cast<T (*)(Attribute)>(castFn))
       );
     }
     template <typename T, typename = DerivedAttrValueCheckT<T>>
     std::optional<DerivedAttrValueIterator<T>> try_value_begin() const {

mlir/include/mlir/IR/BuiltinAttributes.h

 
 //===----------------------------------------------------------------------===//
 // SparseElementsAttr
 //===----------------------------------------------------------------------===//
 
 template <typename T>
 auto SparseElementsAttr::try_value_begin_impl(OverloadToken<T>) const
     -> FailureOr<iterator<T>> {
-  auto zeroValue = getZeroValue<T>();
   auto valueIt = getValues().try_value_begin<T>();
   if (failed(valueIt))
     return failure();
+
+  auto zeroValue = getZeroValue<T>();
   const std::vector<ptrdiff_t> flatSparseIndices(getFlattenedSparseIndices());
   std::function<T(ptrdiff_t)> mapFn =
       [flatSparseIndices{flatSparseIndices}, valueIt{std::move(*valueIt)},
        zeroValue{std::move(zeroValue)}](ptrdiff_t index) {
         // Try to map the current index to one of the sparse indices.
         for (unsigned i = 0, e = flatSparseIndices.size(); i != e; ++i)
           if (flatSparseIndices[i] == index)
             return *std::next(valueIt, i);

mlir/test/IR/elements-attr-interface.mlir

 // RUN: mlir-opt %s -test-elements-attr-interface -verify-diagnostics
 
 // Parsing external resources does not work on big-endian platforms currently
 // XFAIL: target=s390x-{{.*}}
 
 // This test contains various `ElementsAttr` attributes, and tests the support
 // for iterating the values of these attributes using various native C++ types.
 // This tests that the abstract iteration of ElementsAttr works properly, and
 // is properly failable when necessary.
 
 // expected-error@below {{Test iterating `int64_t`: unable to iterate type}}
 // expected-error@below {{Test iterating `uint64_t`: 10, 11, 12, 13, 14}}
 // expected-error@below {{Test iterating `APInt`: 10, 11, 12, 13, 14}}
 // expected-error@below {{Test iterating `IntegerAttr`: 10 : i64, 11 : i64, 12 : i64, 13 : i64, 14 : i64}}
 arith.constant #test.i64_elements<[10, 11, 12, 13, 14]> : tensor<5xi64>
 
+// expected-error@below {{Test iterating `int64_t`: unable to iterate type}}
+// expected-error@below {{Test iterating `uint64_t`: unable to iterate type}}
+// expected-error@below {{Test iterating `APInt`: unable to iterate type}}
+// expected-error@below {{Test iterating `IntegerAttr`: unable to iterate type}}
+arith.constant dense<[1.0, 2.0, 3.0, 4.0, 5.0]> : tensor<5xf32>
+
 // expected-error@below {{Test iterating `int64_t`: 10, 11, 12, 13, 14}}
 // expected-error@below {{Test iterating `uint64_t`: 10, 11, 12, 13, 14}}
 // expected-error@below {{Test iterating `APInt`: 10, 11, 12, 13, 14}}
 // expected-error@below {{Test iterating `IntegerAttr`: 10 : i64, 11 : i64, 12 : i64, 13 : i64, 14 : i64}}
 arith.constant dense<[10, 11, 12, 13, 14]> : tensor<5xi64>
 
 // Check that we don't crash on empty element attributes.
 // expected-error@below {{Test iterating `int64_t`: }}
 // expected-error@below {{Test iterating `uint64_t`: }}
 // expected-error@below {{Test iterating `APInt`: }}
 // expected-error@below {{Test iterating `IntegerAttr`: }}
 arith.constant dense<> : tensor<0xi64>
 
 // Check that we handle an external constant parsed from the config.
 // expected-error@below {{Test iterating `int64_t`: unable to iterate type}}
 // expected-error@below {{Test iterating `uint64_t`: 1, 2, 3}}
 // expected-error@below {{Test iterating `APInt`: unable to iterate type}}
 // expected-error@below {{Test iterating `IntegerAttr`: unable to iterate type}}
 arith.constant #test.e1di64_elements<blob1> : tensor<3xi64>
 
 {-#
   dialect_resources: {
     test: {
       blob1: "0x08000000010000000000000002000000000000000300000000000000"
     }
   }
 #-}