diff --git a/mlir/include/mlir/Dialect/StandardOps/IR/Ops.td b/mlir/include/mlir/Dialect/StandardOps/IR/Ops.td --- a/mlir/include/mlir/Dialect/StandardOps/IR/Ops.td +++ b/mlir/include/mlir/Dialect/StandardOps/IR/Ops.td @@ -455,6 +455,21 @@ let assemblyFormat = "$result attr-dict `:` type($result)"; } +//===----------------------------------------------------------------------===// +// BitcastOp +//===----------------------------------------------------------------------===// + +def BitcastOp : ArithmeticCastOp<"bitcast">, Arguments<(ins AnyType:$in)> { + let summary = "bitcast between values of equal bit width"; + let description = [{ + Bitcast an integer or floating point value to an integer or floating point + value of equal bit width. At runtime this is always a noop cast. + }]; + let hasFolder = 1; + let hasCanonicalizer = 1; +} + + //===----------------------------------------------------------------------===// // BranchOp //===----------------------------------------------------------------------===// diff --git a/mlir/lib/Dialect/StandardOps/IR/Ops.cpp b/mlir/lib/Dialect/StandardOps/IR/Ops.cpp --- a/mlir/lib/Dialect/StandardOps/IR/Ops.cpp +++ b/mlir/lib/Dialect/StandardOps/IR/Ops.cpp @@ -23,6 +23,7 @@ #include "mlir/IR/Value.h" #include "mlir/Support/MathExtras.h" #include "mlir/Transforms/InliningUtils.h" +#include "llvm/ADT/APFloat.h" #include "llvm/ADT/STLExtras.h" #include "llvm/ADT/StringSwitch.h" #include "llvm/Support/FormatVariadic.h" @@ -482,6 +483,69 @@ return success(); } +//===----------------------------------------------------------------------===// +// BitcastOp +//===----------------------------------------------------------------------===// + +bool BitcastOp::areCastCompatible(TypeRange inputs, TypeRange outputs) { + if (inputs.size() != 1 || outputs.size() != 1) + return false; + Type a = inputs.front(), b = outputs.front(); + if (a.isSignlessIntOrFloat() && b.isSignlessIntOrFloat()) + return a.getIntOrFloatBitWidth() == b.getIntOrFloatBitWidth(); + return areVectorCastSimpleCompatible(a, b, areCastCompatible); +} + +OpFoldResult BitcastOp::fold(ArrayRef operands) { + assert(operands.size() == 1 && "bitcastop expects 1 operand"); + auto operand = operands[0]; + if (!operand) + return {}; + + Type resType = getResult().getType(); + + if (auto denseAttr = operand.dyn_cast()) { + Type elType = getElementTypeOrSelf(resType); + return denseAttr.mapValues( + elType, [](const APFloat &f) { return f.bitcastToAPInt(); }); + } + if (auto denseAttr = operand.dyn_cast()) { + Type elType = getElementTypeOrSelf(resType); + return denseAttr.mapValues(elType, [](const APInt &i) { return i; }); + } + + if (auto resIntType = resType.dyn_cast()) { + if (auto floatAttr = operand.dyn_cast()) + return IntegerAttr::get(resType, floatAttr.getValue().bitcastToAPInt()); + if (auto intAttr = operand.dyn_cast()) + return IntegerAttr::get(resType, intAttr.getValue()); + } + if (auto resFloatType = resType.dyn_cast()) { + if (auto intAttr = operand.dyn_cast()) + return FloatAttr::get(resType, APFloat(resFloatType.getFloatSemantics(), + intAttr.getValue())); + if (auto floatAttr = operand.dyn_cast()) + return FloatAttr::get(resType, floatAttr.getValue()); + } + return {}; +} + +static LogicalResult simplifyBitcastOfBitcast(BitcastOp op, + PatternRewriter &rewriter) { + auto *sourceOp = op.getOperand().getDefiningOp(); + if (auto sourceBitcast = dyn_cast_or_null(sourceOp)) { + rewriter.replaceOpWithNewOp(op, op.getResult().getType(), + sourceBitcast.getOperand()); + return success(); + } + return failure(); +} + +void BitcastOp::getCanonicalizationPatterns(RewritePatternSet &results, + MLIRContext *context) { + results.add(&simplifyBitcastOfBitcast); +} + //===----------------------------------------------------------------------===// // BranchOp //===----------------------------------------------------------------------===// diff --git a/mlir/test/Dialect/Standard/canonicalize.mlir b/mlir/test/Dialect/Standard/canonicalize.mlir --- a/mlir/test/Dialect/Standard/canonicalize.mlir +++ b/mlir/test/Dialect/Standard/canonicalize.mlir @@ -331,3 +331,91 @@ %res = select %arg0, %false, %true : i1 return %res : i1 } + +// ----- + +// CHECK-LABEL: @sameTypeBitcast( +// CHECK-SAME: %[[ARG:[a-zA-Z0-9_]*]] +func @sameTypeBitcast(%arg : f32) -> f32 { + // CHECK: return %[[ARG]] + %res = bitcast %arg : f32 to f32 + return %res : f32 +} + +// ----- + +// CHECK-LABEL: @bitcastConstantFPtoI( +func @bitcastConstantFPtoI() -> i32 { + // CHECK: %[[C0:.+]] = constant 0 : i32 + // CHECK: return %[[C0]] + %c0 = constant 0.0 : f32 + %res = bitcast %c0 : f32 to i32 + return %res : i32 +} + +// ----- + +// CHECK-LABEL: @bitcastConstantFPtoI( +func @bitcastConstantFPtoI() -> f32 { + // CHECK: %[[C0:.+]] = constant 0 + // CHECK: return %[[C0]] + %c0 = constant 0 : i32 + %res = bitcast %c0 : i32 to f32 + return %res : f32 +} + +// ----- + +// CHECK-LABEL: @bitcastConstantFPtoI( +func @bitcastConstantFPtoI() -> f32 { + // CHECK: %[[C0:.+]] = constant 0 + // CHECK: return %[[C0]] + %c0 = constant 0 : i32 + %res = bitcast %c0 : i32 to f32 + return %res : f32 +} + +// ----- + +// CHECK-LABEL: @bitcastConstantVecFPtoI( +func @bitcastConstantVecFPtoI() -> vector<3xf32> { + // CHECK: %[[C0:.+]] = constant dense<0 + // CHECK: return %[[C0]] + %c0 = constant dense<0> : vector<3xi32> + %res = bitcast %c0 : vector<3xi32> to vector<3xf32> + return %res : vector<3xf32> +} + +// ----- + +// CHECK-LABEL: @bitcastConstantVecItoFP( +func @bitcastConstantVecItoFP() -> vector<3xi32> { + // CHECK: %[[C0:.+]] = constant dense<0> + // CHECK: return %[[C0]] + %c0 = constant dense<0.0> : vector<3xf32> + %res = bitcast %c0 : vector<3xf32> to vector<3xi32> + return %res : vector<3xi32> +} + +// ----- + +// CHECK-LABEL: @bitcastBackAndForth( +// CHECK-SAME: %[[ARG:[a-zA-Z0-9_]*]] +func @bitcastBackAndForth(%arg : i32) -> i32 { + // CHECK: return %[[ARG]] + %f = bitcast %arg : i32 to f32 + %res = bitcast %f : f32 to i32 + return %res : i32 +} + +// ----- + +// CHECK-LABEL: @bitcastOfBitcast( +// CHECK-SAME: %[[ARG:[a-zA-Z0-9_]*]] +func @bitcastOfBitcast(%arg : i16) -> i16 { + // CHECK: return %[[ARG]] + %f = bitcast %arg : i16 to f16 + %bf = bitcast %f : f16 to bf16 + %res = bitcast %bf : bf16 to i16 + return %res : i16 +} diff --git a/mlir/test/Dialect/Standard/invalid.mlir b/mlir/test/Dialect/Standard/invalid.mlir --- a/mlir/test/Dialect/Standard/invalid.mlir +++ b/mlir/test/Dialect/Standard/invalid.mlir @@ -69,3 +69,11 @@ %0 = constant [1.0 : f32, -1.0 : f64] : complex return } + +// ----- + +func @bitcast_different_bit_widths(%arg : f16) -> f32 { + // expected-error@+1 {{are cast incompatible}} + %res = bitcast %arg : f16 to f32 + return %res : f32 +} diff --git a/mlir/test/Dialect/Standard/ops.mlir b/mlir/test/Dialect/Standard/ops.mlir --- a/mlir/test/Dialect/Standard/ops.mlir +++ b/mlir/test/Dialect/Standard/ops.mlir @@ -80,3 +80,9 @@ %result = constant [0.1 : f64, -1.0 : f64] : complex return %result : complex } + +// CHECK-LABEL: func @bitcast( +func @bitcast(%arg : f32) -> i32 { + %res = bitcast %arg : f32 to i32 + return %res : i32 +}