[mlir][DenseElementsAttr] Allow for custom floating point types in `getValues`

[rriddle]

Summary

Some users have native c++ data types that correspond to the floating point values stored within a DenseElementsAttr that do not have a corresponding native C++ data type(e.g. bfloat16/half/etc). This revision allows for such users to use those native types directly, and removes the need to go through APFloat when the much faster native value path is available.

mlir/include/mlir/IR/Attributes.h

 } // namespace detail
 
 /// An attribute that represents a reference to a dense vector or tensor object.
 ///
 class DenseElementsAttr : public ElementsAttr {
 public:
   using ElementsAttr::ElementsAttr;
 
+  /// Type trait used to check if the given type T is a potentially valid C++
+  /// floating point type that can be used to access the underlying element
+  /// types of a DenseElementsAttr.
+  // TODO: Use std::disjunction when C++17 is supported.
+  template <typename T> struct is_valid_cpp_fp_type {
+    /// The type is a valid floating point type if it is a builtin floating
+    /// point type, or is a potentially user defined floating point type. The
+    /// latter allows for supporting users that have custom types defined for
+    /// bfloat16/half/etc.
+    static inline constexpr bool value =
+        llvm::is_one_of<T, float, double>::value ||
+        (std::numeric_limits<T>::is_specialized &&
+         !std::numeric_limits<T>::is_integer);
+  };
+
   /// Method for support type inquiry through isa, cast and dyn_cast.
   static bool classof(Attribute attr);
 
   /// 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. 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
   /// or floating-point values. Each value is expected to be the same bitwidth
   /// of the element type of 'type'. 'type' must be a vector or tensor with
   /// static shape.
   template <typename T, typename = typename std::enable_if<
                             std::numeric_limits<T>::is_integer ||
-                            llvm::is_one_of<T, float, double>::value>::type>
+                            is_valid_cpp_fp_type<T>::value>::type>
   static DenseElementsAttr get(const ShapedType &type, ArrayRef<T> values) {
     const char *data = reinterpret_cast<const char *>(values.data());
     return getRawIntOrFloat(
         type, ArrayRef<char>(data, values.size() * sizeof(T)), sizeof(T),
         std::numeric_limits<T>::is_integer, std::numeric_limits<T>::is_signed);
   }
 
   /// Constructs a dense integer elements attribute from a single element.
   template <typename T, typename = typename std::enable_if<
                             std::numeric_limits<T>::is_integer ||
-                            llvm::is_one_of<T, float, double>::value ||
+                            is_valid_cpp_fp_type<T>::value ||
                             detail::is_complex_t<T>::value>::type>
   static DenseElementsAttr get(const ShapedType &type, T value) {
     return get(type, llvm::makeArrayRef(value));
   }
 
   /// Constructs a dense complex elements attribute from an array of complex
   /// values. Each value is expected to be the same bitwidth of the element type
   /// of 'type'. 'type' must be a vector or tensor with static shape.
   template <typename T, typename ElementT = typename T::value_type,
             typename = typename std::enable_if<
                 detail::is_complex_t<T>::value &&
                 (std::numeric_limits<ElementT>::is_integer ||
-                 llvm::is_one_of<ElementT, float, double>::value)>::type>
+                 is_valid_cpp_fp_type<ElementT>::value)>::type>
   static DenseElementsAttr get(const ShapedType &type, ArrayRef<T> values) {
     const char *data = reinterpret_cast<const char *>(values.data());
     return getRawComplex(type, ArrayRef<char>(data, values.size() * sizeof(T)),
                          sizeof(T), std::numeric_limits<ElementT>::is_integer,
                          std::numeric_limits<ElementT>::is_signed);
   }
 
   /// Overload of the above 'get' method that is specialized for boolean values.
template <typename T> T getValue(ArrayRef<uint64_t> index) const {
     return *std::next(getValues<T>().begin(), getFlattenedIndex(index));
   }
 
   /// Return the held element values as a range of integer or floating-point
   /// values.
   template <typename T, typename = typename std::enable_if<
                             (!std::is_same<T, bool>::value &&
                              std::numeric_limits<T>::is_integer) ||
-                            llvm::is_one_of<T, float, double>::value>::type>
+                            is_valid_cpp_fp_type<T>::value>::type>
   llvm::iterator_range<ElementIterator<T>> getValues() const {
     assert(isValidIntOrFloat(sizeof(T), std::numeric_limits<T>::is_integer,
                              std::numeric_limits<T>::is_signed));
     const char *rawData = getRawData().data();
     bool splat = isSplat();
     return {ElementIterator<T>(rawData, splat, 0),
             ElementIterator<T>(rawData, splat, getNumElements())};
   }
 
   /// Return the held element values as a range of std::complex.
   template <typename T, typename ElementT = typename T::value_type,
             typename = typename std::enable_if<
                 detail::is_complex_t<T>::value &&
                 (std::numeric_limits<ElementT>::is_integer ||
-                 llvm::is_one_of<ElementT, float, double>::value)>::type>
+                 is_valid_cpp_fp_type<ElementT>::value)>::type>
   llvm::iterator_range<ElementIterator<T>> getValues() const {
     assert(isValidComplex(sizeof(T), std::numeric_limits<ElementT>::is_integer,
                           std::numeric_limits<ElementT>::is_signed));
     const char *rawData = getRawData().data();
     bool splat = isSplat();
     return {ElementIterator<T>(rawData, splat, 0),
             ElementIterator<T>(rawData, splat, getNumElements())};
   }
getZeroValue() const {
     APFloat floatZero = getZeroAPFloat();
     return {floatZero, floatZero};
   }
 
   /// Get a zero for an C++ integer, float, StringRef, or complex type.
   template <typename T>
   typename std::enable_if<
       std::numeric_limits<T>::is_integer ||
-          llvm::is_one_of<T, float, double, StringRef>::value ||
+          DenseElementsAttr::is_valid_cpp_fp_type<T>::value ||
+          std::is_same<T, StringRef>::value ||
           (detail::is_complex_t<T>::value &&
            !llvm::is_one_of<T, std::complex<APInt>,
                             std::complex<APFloat>>::value),
       T>::type
   getZeroValue() const {
     return T();
   }