[mlir][complex] Add a `complex.bitcast` operation

[rsuderman]

Summary

Converting between a complex<f32> to i64 could be useful for handling interop
between the arith and complex dialects.

[rsuderman]

Rebased on main

[jpienaar]

mlir/lib/Dialect/Complex/IR/ComplexOps.cpp

91

Elide trivial braces

95

Why is this allowed vs keeping this purely on complex?

150

So this could result in changing arith Bitcast into a complex one, but only if operand is result of bitcast?

[rsuderman]

Updated for jpienaar@ comments. Included adding complex.bitcast -> arith.bitcast canonicalizer.

[rsuderman]

mlir/lib/Dialect/Complex/IR/ComplexOps.cpp

95

We want to allow bitcasting between i64 and complex<f32> so both need to be allowed. However we want to avoid complex.bitcast being used between unsupported types (e.g. tensor<complex<f32>>) for now.

150

Essentially yes. If we see a complex.bitcast -> arith.bitcast we merge this into a single complex.bitcast. I would argue its just sucking the arith.bitcast into the complex.bitcast but thats just a minor detail.

I realized I should include a complex.bitcast to arith.bitcast conversion so that if two complex.bitcasts are merged such that they are a valid arith.bitcast, we should just canonocalize to it.

mlir/include/mlir/Dialect/Complex/IR/ComplexOps.td

let description = [{
     Example:
 
     ```mlir
     %a = complex.atan2 %b, %c : complex<f32>
     ```
   }];
 }
 
+
+//===----------------------------------------------------------------------===//
+// Bitcast
+//===----------------------------------------------------------------------===//
+
+def BitcastOp : Complex_Op<"bitcast", [Pure]> {
+
+  let summary = "computes bitcast between complex and equal arith types";
+  let description = [{
+
+    Example:
+
+    ```mlir
+         %a = complex.bitcast %b : complex<f32> -> i64
+    ```
+  }];
+  let assemblyFormat = "$operand attr-dict `:` type($operand) `to` type($result)";
+  let arguments = (ins AnyType:$operand);
+  let results = (outs AnyType:$result);
+
+  let hasCanonicalizer = 1;
+  let hasFolder = 1;
+  let hasVerifier = 1;
+}
+
 //===----------------------------------------------------------------------===//
 // ConstantOp
 //===----------------------------------------------------------------------===//
 
 def ConstantOp : Complex_Op<"constant", [
     ConstantLike, Pure,
     DeclareOpInterfaceMethods<OpAsmOpInterface, ["getAsmResultNames"]>
   ]> {

mlir/lib/Dialect/Complex/IR/ComplexOps.cpp

return emitOpError()
            << im.getType()
            << ") to match the element type of the op's return type ("
            << complexEltTy << ")";
   }
   return success();
 }
 
 //===----------------------------------------------------------------------===//
+// BitcastOp
+//===----------------------------------------------------------------------===//
+
+OpFoldResult BitcastOp::fold(FoldAdaptor bitcast) {
+  if (getOperand().getType() == getType())
+    return getOperand();
+
+  return {};
+}
+
+LogicalResult BitcastOp::verify() {
+  auto operandType = getOperand().getType();
+  auto resultType = getType();
+
+  // We allow this to be legal as it can be folded away.
+  if (operandType == resultType)
+    return success();
+
+  if (!operandType.isIntOrFloat() && !isa<ComplexType>(operandType)) {
+    return emitOpError("operand must be int/float/complex");
+  }
+
+  if (!resultType.isIntOrFloat() && !isa<ComplexType>(resultType)) {
+    return emitOpError("result must be int/float/complex");
+  }
+
+  if (isa<ComplexType>(operandType) == isa<ComplexType>(resultType)) {
+    return emitOpError("requires input or output is a complex type");
+  }
+
+  if (isa<ComplexType>(resultType))
+    std::swap(operandType, resultType);
+
+  int32_t operandBitwidth = dyn_cast<ComplexType>(operandType)
+                                .getElementType()
+                                .getIntOrFloatBitWidth() *
+                            2;
+  int32_t resultBitwidth = resultType.getIntOrFloatBitWidth();
+
+  if (operandBitwidth != resultBitwidth) {
+    return emitOpError("casting bitwidths do not match");
+  }
+
+  return success();
+}
+
+struct MergeComplexBitcast final : OpRewritePattern<BitcastOp> {
+  using OpRewritePattern<BitcastOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(BitcastOp op,
+                                PatternRewriter &rewriter) const override {
+    if (auto defining = op.getOperand().getDefiningOp<BitcastOp>()) {
+      rewriter.replaceOpWithNewOp<BitcastOp>(op, op.getType(),
+                                             defining.getOperand());
+      return success();
+    }
+
+    if (auto defining = op.getOperand().getDefiningOp<arith::BitcastOp>()) {
+      rewriter.replaceOpWithNewOp<BitcastOp>(op, op.getType(),
+                                             defining.getOperand());
+      return success();
+    }
+
+    return failure();
+  }
+};
+
+struct MergeArithBitcast final : OpRewritePattern<arith::BitcastOp> {
+  using OpRewritePattern<arith::BitcastOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(arith::BitcastOp op,
+                                PatternRewriter &rewriter) const override {
+    if (auto defining = op.getOperand().getDefiningOp<complex::BitcastOp>()) {
+      rewriter.replaceOpWithNewOp<complex::BitcastOp>(op, op.getType(),
+                                                      defining.getOperand());
+      return success();
+    }
+
+    return failure();
+  }
+};
+
+struct ArithBitcast final : OpRewritePattern<BitcastOp> {
+  using OpRewritePattern<complex::BitcastOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(BitcastOp op,
+                                PatternRewriter &rewriter) const override {
+    if (isa<ComplexType>(op.getType()) ||
+        isa<ComplexType>(op.getOperand().getType()))
+      return failure();
+
+    rewriter.replaceOpWithNewOp<arith::BitcastOp>(op, op.getType(),
+                                                  op.getOperand());
+    return success();
+  }
+};
+
+void BitcastOp::getCanonicalizationPatterns(RewritePatternSet &results,
+                                            MLIRContext *context) {
+  results.add<ArithBitcast, MergeComplexBitcast, MergeArithBitcast>(context);
+}
+
+//===----------------------------------------------------------------------===//
 // CreateOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult CreateOp::fold(FoldAdaptor adaptor) {
   // Fold complex.create(complex.re(op), complex.im(op)).
   if (auto reOp = getOperand(0).getDefiningOp<ReOp>()) {
     if (auto imOp = getOperand(1).getDefiningOp<ImOp>()) {
       if (reOp.getOperand() == imOp.getOperand()) {

mlir/test/Dialect/Complex/canonicalize.mlir

 func.func @mul_one_f128(%arg0: f128, %arg1: f128) -> complex<f128> {
   %create = complex.create %arg0, %arg1: complex<f128>  
   %one = complex.constant [1.0 : f128, 0.0 : f128] : complex<f128>
   %mul = complex.mul %create, %one : complex<f128>
   // CHECK: %[[CREATE:.*]] = complex.create %[[ARG0]], %[[ARG1]] : complex<f128>
   // CHECK-NEXT: return %[[CREATE]]
   return %mul : complex<f128>
 }
+
+// CHECK-LABEL: func @fold_between_complex
+//  CHECK-SAME: %[[ARG0:.*]]: complex<f32>
+func.func @fold_between_complex(%arg0 : complex<f32>) -> complex<f32> {
+  %0 = complex.bitcast %arg0 : complex<f32> to i64
+  %1 = complex.bitcast %0 : i64 to complex<f32>
+  // CHECK: return %[[ARG0]] : complex<f32>
+  func.return %1 : complex<f32>
+}
+
+// CHECK-LABEL: func @fold_between_i64
+//  CHECK-SAME: %[[ARG0:.*]]: i64
+func.func @fold_between_i64(%arg0 : i64) -> i64 {
+  %0 = complex.bitcast %arg0 : i64 to complex<f32>
+  %1 = complex.bitcast %0 : complex<f32> to i64
+  // CHECK: return %[[ARG0]] : i64
+  func.return %1 : i64
+}
+
+// CHECK-LABEL: func @canon_arith_bitcast
+//  CHECK-SAME: %[[ARG0:.*]]: f64
+func.func @canon_arith_bitcast(%arg0 : f64) -> i64 {
+  %0 = complex.bitcast %arg0 : f64 to complex<f32>
+  %1 = complex.bitcast %0 : complex<f32> to i64
+  // CHECK: %[[R0:.+]] = arith.bitcast %[[ARG0]]
+  // CHECK: return %[[R0]] : i64
+  func.return %1 : i64
+}
+
+
+// CHECK-LABEL: func @double_bitcast
+//  CHECK-SAME: %[[ARG0:.*]]: f64
+func.func @double_bitcast(%arg0 : f64) -> complex<f32> {
+  // CHECK: %[[R0:.+]] = complex.bitcast %[[ARG0]]
+  %0 = arith.bitcast %arg0 : f64 to i64
+  %1 = complex.bitcast %0 : i64 to complex<f32>
+  // CHECK: return %[[R0]] : complex<f32>
+  func.return %1 : complex<f32>
+}
+
+// CHECK-LABEL: func @double_reverse_bitcast
+//  CHECK-SAME: %[[ARG0:.*]]: complex<f32>
+func.func @double_reverse_bitcast(%arg0 : complex<f32>) -> f64 {
+  // CHECK: %[[R0:.+]] = complex.bitcast %[[ARG0]]
+  %0 = complex.bitcast %arg0 : complex<f32> to i64
+  %1 = arith.bitcast %0 : i64 to f64
+  // CHECK: return %[[R0]] : f64
+  func.return %1 : f64
+}

mlir/test/Dialect/Complex/invalid.mlir

 
 // -----
 
 func.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 = complex.constant [1.0 : f32, -1.0 : f64] : complex<f64>
   return
 }
+
+// -----
+
+func.func @complex_bitcast_i64(%arg0 : i64) {
+  // expected-error @+1 {{op requires input or output is a complex type}}
+  %0 = complex.bitcast %arg0: i64 to f64
+  return
+}

mlir/test/Dialect/Complex/ops.mlir

func.func @ops(%f: f32) {
   %rsqrt = complex.rsqrt %complex : complex<f32>
 
   // CHECK: complex.atan2 %[[C]], %[[C]] : complex<f32>
   %atan2 = complex.atan2 %complex, %complex : complex<f32>
 
   // CHECK: complex.tan %[[C]] : complex<f32>
   %tan = complex.tan %complex : complex<f32>
 
+  // CHECK: complex.bitcast %[[C]]
+  %i64 = complex.bitcast %complex : complex<f32> to i64
+
   return
 }