[Mlir-commits] [mlir] e2c97d4 - [MLIR] Add a bitcast method to DenseElementsAttr

[Geoffrey Martin-Noble]

Author: Geoffrey Martin-Noble
Date: 2021-08-16T17:13:35-07:00
New Revision: e2c97d4484468a498def2527136edeb5abc332d9

URL: https://github.com/llvm/llvm-project/commit/e2c97d4484468a498def2527136edeb5abc332d9
DIFF: https://github.com/llvm/llvm-project/commit/e2c97d4484468a498def2527136edeb5abc332d9.diff

LOG: [MLIR] Add a bitcast method to DenseElementsAttr

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

Reviewed By: rriddle

Differential Revision: https://reviews.llvm.org/D107612

Added: 
    

Modified: 
    mlir/include/mlir/IR/BuiltinAttributes.h
    mlir/lib/Dialect/StandardOps/IR/Ops.cpp
    mlir/lib/IR/BuiltinAttributes.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/IR/BuiltinAttributes.h b/mlir/include/mlir/IR/BuiltinAttributes.h
index 28ced2deb86f0..c71c5d4bbd4ae 100644
--- a/mlir/include/mlir/IR/BuiltinAttributes.h
+++ b/mlir/include/mlir/IR/BuiltinAttributes.h
@@ -560,6 +560,11 @@ class DenseElementsAttr : public ElementsAttr {
   /// 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.

diff  --git a/mlir/lib/Dialect/StandardOps/IR/Ops.cpp b/mlir/lib/Dialect/StandardOps/IR/Ops.cpp
index c0d18193b2608..0c4b1441dd4cd 100644
--- a/mlir/lib/Dialect/StandardOps/IR/Ops.cpp
+++ b/mlir/lib/Dialect/StandardOps/IR/Ops.cpp
@@ -512,16 +512,8 @@ OpFoldResult BitcastOp::fold(ArrayRef<Attribute> operands) {
 
   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>())

diff  --git a/mlir/lib/IR/BuiltinAttributes.cpp b/mlir/lib/IR/BuiltinAttributes.cpp
index 85914be243b92..5a754dfe7c3f5 100644
--- a/mlir/lib/IR/BuiltinAttributes.cpp
+++ b/mlir/lib/IR/BuiltinAttributes.cpp
@@ -1024,7 +1024,6 @@ DenseElementsAttr DenseElementsAttr::reshape(ShapedType newType) {
   if (curType == newType)
     return *this;
 
-  (void)curType;
   assert(newType.getElementType() == curType.getElementType() &&
          "expected the same element type");
   assert(newType.getNumElements() == curType.getNumElements() &&
@@ -1032,6 +1031,23 @@ DenseElementsAttr DenseElementsAttr::reshape(ShapedType newType) {
   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 {