[MLIR] Add a bitcast method to DenseElementsAttr

[GMNGeoffrey]

Summary

This method bitcasts a DenseElementsAttr elementwise to one of the same
shape with a different element type.

[GMNGeoffrey]

This modifies the bitcast folders added in https://reviews.llvm.org/D105376, so is dependent on that patch.

[rriddle]

mlir/lib/IR/BuiltinAttributes.cpp

1027

This should be true on construction.

1044

Is this necessary? It's used above.

1046–1047

If we use getDenseElementBitWidth here, you could drop the isIntOrFloat checks and ideally support complex as well.

1049

You'll should also probably assert that the number of elements match.

mlir/lib/IR/BuiltinTypes.cpp

448–455 ↗

(On Diff #364649)

This looks unrelated.

[GMNGeoffrey]

• Respond to review comments

[GMNGeoffrey]

mlir/lib/IR/BuiltinAttributes.cpp

1027

It was intended to be a check for integer/float. Deleted though, per your comment below.

1044

It's not. I was following the pattern in reshape, which I assumed was intended as defense against some future mistake (or just a general principle of putting void casts before asserts with anything they use. I'll delete this and the one above, if this isn't some pattern I'm supposed to be following.

1046–1047

Done.

1049

That's true by construction with clone (this accepts an element type not a ShapedType)

mlir/lib/IR/BuiltinTypes.cpp

448–455 ↗

(On Diff #364649)

This was something I added while debugging. I've split it into a separate patch: https://reviews.llvm.org/D107676. Better debug messages are generally helpful though, I think :-)

mlir/include/mlir/IR/BuiltinAttributes.h

public:
   // Mutation Utilities
   //===--------------------------------------------------------------------===//
 
   /// Return a new DenseElementsAttr that has the same data as the current
   /// attribute, but has been reshaped to 'newType'. The new type must have the
   /// same total number of elements as well as element type.
   DenseElementsAttr reshape(ShapedType newType);
 
+  /// Return a new DenseElementsAttr that has the same data as the current
+  /// attribute, but has bitcast elements to 'newElType'. The new type must have
+  /// the same bitwidth as the current element type.
+  DenseElementsAttr bitcast(Type newElType);
+
   /// Generates a new DenseElementsAttr by mapping each int value to a new
   /// underlying APInt. The new values can represent either an integer or float.
   /// This underlying type must be an DenseIntElementsAttr.
   DenseElementsAttr mapValues(Type newElementType,
                               function_ref<APInt(const APInt &)> mapping) const;
 
   /// Generates a new DenseElementsAttr by mapping each float value to a new
   /// underlying APInt. the new values can represent either an integer or float.

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

OpFoldResult BitcastOp::fold(ArrayRef<Attribute> operands) {
   }
 
   auto operand = operands[0];
   if (!operand)
     return {};
 
   Type resType = getResult().getType();
 
-  if (auto denseAttr = operand.dyn_cast<DenseFPElementsAttr>()) {
-    Type elType = getElementTypeOrSelf(resType);
-    return denseAttr.mapValues(
-        elType, [](const APFloat &f) { return f.bitcastToAPInt(); });
-  }
-  if (auto denseAttr = operand.dyn_cast<DenseIntElementsAttr>()) {
-    Type elType = getElementTypeOrSelf(resType);
-    // mapValues does its own bitcast to the target type.
-    return denseAttr.mapValues(elType, [](const APInt &i) { return i; });
-  }
+  if (auto denseAttr = operand.dyn_cast<DenseElementsAttr>())
+    return denseAttr.bitcast(resType.cast<ShapedType>().getElementType());
 
   APInt bits;
   if (auto floatAttr = operand.dyn_cast<FloatAttr>())
     bits = floatAttr.getValue().bitcastToAPInt();
   else if (auto intAttr = operand.dyn_cast<IntegerAttr>())
     bits = intAttr.getValue();
   else
     return {};

mlir/lib/IR/BuiltinAttributes.cpp

 /// Return a new DenseElementsAttr that has the same data as the current
 /// attribute, but has been reshaped to 'newType'. The new type must have the
 /// same total number of elements as well as element type.
 DenseElementsAttr DenseElementsAttr::reshape(ShapedType newType) {
   ShapedType curType = getType();
   if (curType == newType)
     return *this;
 
-  (void)curType;
   assert(newType.getElementType() == curType.getElementType() &&
          "expected the same element type");
   assert(newType.getNumElements() == curType.getNumElements() &&
          "expected the same number of elements");
   return DenseIntOrFPElementsAttr::getRaw(newType, getRawData(), isSplat());
 }
 
+/// Return a new DenseElementsAttr that has the same data as the current
+/// attribute, but has bitcast elements such that it is now 'newType'. The new
+/// type must have the same shape and element types of the same bitwidth as the
+/// current type.
+DenseElementsAttr DenseElementsAttr::bitcast(Type newElType) {
+  ShapedType curType = getType();
+  Type curElType = curType.getElementType();
+  if (curElType == newElType)
+    return *this;
+
+  assert(getDenseElementBitWidth(newElType) ==
+             getDenseElementBitWidth(curElType) &&
+         "expected element types with the same bitwidth");
+  return DenseIntOrFPElementsAttr::getRaw(curType.clone(newElType),
+                                          getRawData(), isSplat());
+}
+
 DenseElementsAttr
 DenseElementsAttr::mapValues(Type newElementType,
                              function_ref<APInt(const APInt &)> mapping) const {
   return cast<DenseIntElementsAttr>().mapValues(newElementType, mapping);
 }
 
 DenseElementsAttr DenseElementsAttr::mapValues(
     Type newElementType, function_ref<APInt(const APFloat &)> mapping) const {