[MLIR][std] Introduce bitcast operation

[GMNGeoffrey]

Summary

This patch introduces a bitcast operation to the standard dialect.
RFC: https://llvm.discourse.group/t/rfc-introduce-a-bitcast-op/3774

[GMNGeoffrey]

Add folds and canonicalization

[GMNGeoffrey]

inb4 River asks me to remove trivial braces

[mehdi_amini]

LG in general, some comments inline

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

462

The argument is "AnyType" but the description makes me think that if should be IntOrFloat or something like that.

Similarly for the result type, there is no verifier here to check on the "equal bit width"?

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

503

There is another fold here: the identity bitcast can just return its input as-is.

Edit: actually I see a test for it, but I don't know how it works? Is there a base-class / interface for cast operation that takes care of it?

538

Can't we change the operand in-place instead of creating a new operation here?

[mehdi_amini]

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

465

Actually you support vector as well according to the tests.

466

I'm not sure about the noop part: a machine can have different registers for floating point and integer and would require a copy across register files.

[silvas]

Please add lowering to LLVM for this too (doesn't have to be this patch precisely, but a fast follow-on please).

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

466

Yeah, LLVM uses the terminology that it's like storing the value to memory with one type and reading it back in another type. Not sure how we can phrase it.

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

491

you can assert/ignore this. Otherwise the op is verifier-invalid.

499

Is this folder correct? What if target endianness is different -- how will it match the target behavior?

Generally speaking, I would suggest that we follow how LLVM does this, and leave a comment saying that we are attempting to follow the LLVM folding behavior when we have similar information available (LLVM always has datalayout).

517

are all of these code paths covered in the tests?

For example, I don't think I see one that would trigger nontrivial float->float conversion. f16->bf16 might be a case to exercise that.

[GMNGeoffrey]

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

462

I followed the style of the other cast operations here where this is all done in the cast compatible interface method. I'm not sure that that's actually better, but it seems to be what's done here. I could pre-factor or follow up to express more of this in the op definition for all of these ops

465

Yeah again I was following the rest of the ops here which just talk about the scalar aspect. Maybe they should all be updated to talk about vectors as well.

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

503

I believe it's encapsulated in the interface.

517

I've got a cast that uses bf16, but it's testing the cast-of-cast path. I don't think we actually have a way to trigger same-width different integer types, do we?

[mehdi_amini]

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

517

signed integer to unsigned integer maybe?

538

(I don't know if you missed this comment?)

[GMNGeoffrey]

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

517

I think standard ops aren't supposed to operated on the [un]signed integer variants?

538

No just haven't gotten to it yet and didn't have any questions about it :-)

[GMNGeoffrey]

Some pre-factor cleanup: https://reviews.llvm.org/D105934

[GMNGeoffrey]

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

462

Addressed in https://reviews.llvm.org/D105934

465

Addressed in https://reviews.llvm.org/D105934

[GMNGeoffrey]

Add more tests and clean up canonicalizations

[GMNGeoffrey]

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

466

LLVM says

It is always a no-op cast because no bits change with this conversion.

https://llvm.org/docs/LangRef.html#bitcast-to-instruction

Maybe that has some specific definition in LLVM?

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

491

Done

499

The llvm behavior is only different for non-elementwise bitcasts of vectors (which we're not doing here):

There is a caveat for bitcasts involving vector types in relation to endianess. For example bitcast <2 x i8> <value> to i16 puts element zero of the vector in the least significant bits of the i16 for little-endian while element zero ends up in the most significant bits for big-endian.

https://llvm.org/docs/LangRef.html#bitcast-to-instruction

It doesn't reference endianness otherwise. I think *usually* endianness shouldn't matter because the endianness of the source and target types would be the same, but I guess that wouldn't apply on weird architectures where floating-point is big-endian and integer is little-endian.

I think we've got two potential definitions here:

1. The bitcast operation does a target-independent bitcast of the arbitrary-width value. I believe that's what these folders implement, assuming I understood how bitcastToAPInt works. It does mean that the statement that this is always a no-op is not true because a cast from f64 to i64 on a 32-bit system with big-endian floating point and little-endian integer would need to swap the words.
2. The return value is target-dependent. The runtime no-op is still not true because of the float/int registers as Mehdi points out.

Given that we don't have target specifications at this level, I think I like (1) better. Then we have a clearly-defined target-independent operation for the compiler and an operation that is free on non-weird architectures, but needs special handling at the target-specific level for weird architectures.

It looks like the LLVM constant folding only handles casts to vectors: https://github.com/llvm/llvm-project/blob/18c19414eb70578d4c487d6f4b0f438aead71d6a/llvm/lib/Analysis/ConstantFolding.cpp#L149, so we can't really use that as a reference.

The other question would be how to lower to/from such an operation. For instance, I'm interested in lowering mhlo.bitcast, which really doesn't explain its view on endianness in much detail

so machines with different floating-point representations will give different results

https://www.tensorflow.org/xla/operation_semantics#bitcastconverttype

Implementation defined sounds like a free pass there :-P

Lowering to LLVM, would this op be able to be converted directly to an llvm bitcast? I think not, and that LLVM has the second semantics above. When lowering to LLVM, do we have the target information yet? Because then it would be easy to add a special case for the mixed-endian architectures.

538

Done.

[GMNGeoffrey]

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

499

Or maybe not and we go with definition 2 to match LLVM, in which case I think constant folding is in-general impossible

[GMNGeoffrey]

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

499

Ping on this?

[mehdi_amini]

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

499

Is "mixed endian" a thing?

[GMNGeoffrey]

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

499

If by that you mean having different endianness for float and int, apparently yes. There are also other endianness formats that just kind of randomly strew the bytes about, apparently

[GMNGeoffrey]

Ping on this? I can bring the question of exact bitcast semantics back to the RFC if you'd like

[silvas]

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

499

I was referring to "the return value is target dependent" case, which would mean that we could not fold without further information abotu the target. It seems that LLVM has a strong precedent for ignoring that issue, which is equivalent to defining the fp format in terms of the bit positions in the integer format. The critical code for the scalar case in LLVM is here: https://github.com/llvm/llvm-project/blob/18c19414eb70578d4c487d6f4b0f438aead71d6a/llvm/lib/IR/ConstantFold.cpp#L168 I would link to that code and explain that everything we are doing here is consistent with that. The DenseIntElementsAttr case seems like it requires further explanation though along with the DenseFPElementsAttr case.

[GMNGeoffrey]

• Clarify op and folding semantics wrt endianness

[GMNGeoffrey]

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

499

I added a bunch more to operation semantics and some comments to the folding methods

[rriddle]

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

537

You should be able to move this to the folder, updating the root operation is fine you just need to return the existing result of the op for the OpFoldResult.

[GMNGeoffrey]

• Move bitcast-of-bitcast simplification to folder

[GMNGeoffrey]

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

537

Oh nice! Done.

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

let description = [{
     "atomic_yield" yields an SSA value from a GenericAtomicRMWOp region.
   }];
 
   let arguments = (ins AnyType:$result);
   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
 //===----------------------------------------------------------------------===//
 
 def BranchOp : Std_Op<"br",
     [DeclareOpInterfaceMethods<BranchOpInterface, ["getSuccessorForOperands"]>,
      NoSideEffect, Terminator]> {
   let summary = "branch operation";
   let description = [{

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

 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/TypeUtilities.h"
 #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"
 #include "llvm/Support/raw_ostream.h"
 
 #include "mlir/Dialect/StandardOps/IR/OpsDialect.cpp.inc"
 
 // Pull in all enum type definitions and utility function declarations.
static LogicalResult verify(AtomicYieldOp op) {
   Type resultType = op.result().getType();
   if (parentType != resultType)
     return op.emitOpError() << "types mismatch between yield op: " << resultType
                             << " and its parent: " << parentType;
   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<Attribute> 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<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);
+    return denseAttr.mapValues(elType, [](const APInt &i) { return i; });
+  }
+
+  if (auto resIntType = resType.dyn_cast<IntegerType>()) {
+    if (auto floatAttr = operand.dyn_cast<FloatAttr>())
+      return IntegerAttr::get(resType, floatAttr.getValue().bitcastToAPInt());
+    if (auto intAttr = operand.dyn_cast<IntegerAttr>())
+      return IntegerAttr::get(resType, intAttr.getValue());
+  }
+  if (auto resFloatType = resType.dyn_cast<FloatType>()) {
+    if (auto intAttr = operand.dyn_cast<IntegerAttr>())
+      return FloatAttr::get(resType, APFloat(resFloatType.getFloatSemantics(),
+                                             intAttr.getValue()));
+    if (auto floatAttr = operand.dyn_cast<FloatAttr>())
+      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<BitcastOp>(sourceOp)) {
+    rewriter.replaceOpWithNewOp<BitcastOp>(op, op.getResult().getType(),
+                                           sourceBitcast.getOperand());
+    return success();
+  }
+  return failure();
+}
+
+void BitcastOp::getCanonicalizationPatterns(RewritePatternSet &results,
+                                            MLIRContext *context) {
+  results.add(&simplifyBitcastOfBitcast);
+}
+
+//===----------------------------------------------------------------------===//
 // BranchOp
 //===----------------------------------------------------------------------===//
 
 /// Given a successor, try to collapse it to a new destination if it only
 /// contains a passthrough unconditional branch. If the successor is
 /// collapsable, `successor` and `successorOperands` are updated to reference
 /// the new destination and values. `argStorage` is used as storage if operands
 /// to the collapsed successor need to be remapped. It must outlive uses of

mlir/test/Dialect/Standard/canonicalize.mlir

 //       CHECK:       %[[trueval:.+]] = constant true
 //       CHECK:       %{{.+}} = xor %arg0, %[[trueval]] : i1
 func @selToNot(%arg0: i1) -> i1 {
   %true = constant true
   %false = constant false
   %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
+}

mlir/test/Dialect/Standard/invalid.mlir

 
 // -----
 
 func @complex_constant_two_different_element_types() {
   // expected-error @+1 {{requires attribute's element types ('f32', 'f64') to match the element type of the op's return type ('f64')}}
   %0 = constant [1.0 : f32, -1.0 : f64] : complex<f64>
   return
 }
+
+// -----
+
+func @bitcast_different_bit_widths(%arg : f16) -> f32 {
+  // expected-error@+1 {{are cast incompatible}}
+  %res = bitcast %arg : f16 to f32
+  return %res : f32
+}

mlir/test/Dialect/Standard/ops.mlir

func @constant_complex_f32() -> complex<f32> {
   return %result : complex<f32>
 }
 
 // CHECK-LABEL: func @constant_complex_f64(
 func @constant_complex_f64() -> complex<f64> {
   %result = constant [0.1 : f64, -1.0 : f64] : complex<f64>
   return %result : complex<f64>
 }
+
+// CHECK-LABEL: func @bitcast(
+func @bitcast(%arg : f32) -> i32 {
+  %res = bitcast %arg : f32 to i32
+  return %res : i32
+}