[mlir][vector] Add constant folding for fp16 to fp32 bitcast

[antiagainst]

Summary

Depends On D96040

[antiagainst]

Update

[ThomasRaoux]

mlir/lib/Dialect/Vector/VectorOps.cpp

2826

I believe passing the value as float can be risky as Nan encoding can be changed. I'm not positive this is a problem here but it might be better to create an APFloat directly from bits?

[antiagainst]

Address comments

mlir/lib/Dialect/Vector/VectorOps.cpp

2826

Good point! Changed to directly use APFloat.

[antiagainst]

Update

mlir/lib/Dialect/Vector/VectorOps.cpp

 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/TypeUtilities.h"
 #include "mlir/Support/LLVM.h"
 #include "mlir/Support/MathExtras.h"
 #include "llvm/ADT/StringSet.h"
+#include "llvm/ADT/bit.h"
 #include <numeric>
 
 using namespace mlir;
 using namespace mlir::vector;
 
 /// Helper enum to classify mask value.
 enum class MaskFormat {
   AllTrue = 0,
OpFoldResult BitCastOp::fold(ArrayRef<Attribute> operands) {
   if (source().getType() == result().getType())
     return source();
 
   // Canceling bitcasts.
   if (auto otherOp = source().getDefiningOp<BitCastOp>())
     if (result().getType() == otherOp.source().getType())
       return otherOp.source();
 
+  Attribute sourceConstant = operands.front();
+  if (!sourceConstant)
+    return {};
+
+  Type srcElemType = getSourceVectorType().getElementType();
+  Type dstElemType = getResultVectorType().getElementType();
+
+  if (auto floatPack = sourceConstant.dyn_cast<DenseFPElementsAttr>()) {
+    if (floatPack.isSplat()) {
+      auto splat = floatPack.getSplatValue<FloatAttr>();
+
+      // Casting fp16 into fp32.
+      if (srcElemType.isF16() && dstElemType.isF32()) {
+        uint32_t bits = static_cast<uint32_t>(
+            splat.getValue().bitcastToAPInt().getZExtValue());
+        // Duplicate the 16-bit pattern.
+        bits = (bits << 16) | (bits & 0xffff);
+        APInt intBits(32, bits);
+        APFloat floatBits(llvm::APFloat::IEEEsingle(), intBits);
+        return DenseElementsAttr::get(getResultVectorType(), floatBits);
+      }
+    }
+  }
+
   return {};
 }
 
 //===----------------------------------------------------------------------===//
 // TypeCastOp
 //===----------------------------------------------------------------------===//
 
 static SmallVector<int64_t, 8> extractShape(MemRefType memRefType) {

mlir/test/Dialect/Vector/canonicalize.mlir

 //  CHECK:        return %[[A]], %[[B]] : vector<4x8xf32>, vector<2xi32>
 func @bitcast_folding(%I1: vector<4x8xf32>, %I2: vector<2xi32>) -> (vector<4x8xf32>, vector<2xi32>) {
   %0 = vector.bitcast %I1 : vector<4x8xf32> to vector<4x8xf32>
   %1 = vector.bitcast %I2 : vector<2xi32> to vector<4xi16>
   %2 = vector.bitcast %1 : vector<4xi16> to vector<2xi32>
   return %0, %2 : vector<4x8xf32>, vector<2xi32>
 }
 
+// CHECK-LABEL: func @bitcast_f16_to_f32
+//              bit pattern: 0x00000000
+//       CHECK: %[[CST0:.+]] = constant dense<0.000000e+00> : vector<4xf32>
+//              bit pattern: 0x40004000
+//       CHECK: %[[CST1:.+]] = constant dense<2.00390625> : vector<4xf32>
+//       CHECK: return %[[CST0]], %[[CST1]]
+func @bitcast_f16_to_f32() -> (vector<4xf32>, vector<4xf32>) {
+  %cst0 = constant dense<0.0> : vector<8xf16> // bit pattern: 0x0000
+  %cst1 = constant dense<2.0> : vector<8xf16> // bit pattern: 0x4000
+  %cast0 = vector.bitcast %cst0: vector<8xf16> to vector<4xf32>
+  %cast1 = vector.bitcast %cst1: vector<8xf16> to vector<4xf32>
+  return %cast0, %cast1: vector<4xf32>, vector<4xf32>
+}
+
 // -----
 
 // CHECK-LABEL: broadcast_folding1
 //       CHECK: %[[CST:.*]] = constant dense<42> : vector<4xi32>
 //   CHECK-NOT: vector.broadcast
 //       CHECK: return %[[CST]]
 func @broadcast_folding1() -> vector<4xi32> {
   %0 = constant 42 : i32