[mlir][complex] Initial support for FastMath flag when converting to LLVM

[Lewuathe]

Summary

This change contains the initial support of FastMath flag in complex dialect. Similar to what we did in Arith dialect, fastmath attributes in the complex dialect are directly mapped to the corresponding LLVM fastmath flags.

In this diff,

• Definition of FastMathAttr as a custom attribute in the Complex dialect that inherits from the EnumAttr class.
• Definition of ComplexFastMathInterface, which is an interface that is implemented by operations that have a complex::fastmath attribute.
• Declaration of a default-valued fastmath attribute for unary and arithmetic operations in the Complex dialect.
• Conversion code to lower arithmetic fastmath flags to LLVM fastmath flags

NOT in this diff (but planned and progressively implemented):

• Documentation of flag meanings
• Support the fastmath flag conversion to Arith dialect
• Folding/rewrite implementations that are enabled by fastmath flags (although it's the original motivation to support the flag)

RFC: https://discourse.llvm.org/t/rfc-fastmath-flags-support-in-complex-dialect/71981

[kiranchandramohan]

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

29–55 ↗

(On Diff #544295)

Isn't it possible to reuse these definitions from the Arithmetic Dialect (ArithBase.td)?

[kuhar]

+r Rob for complex

[mehdi_amini]

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

29–55 ↗

(On Diff #544295)

+1: I'd like to not see a new attribute here, see https://discourse.llvm.org/t/place-for-fastmath-attributes/69933

[Lewuathe]

Thanks. I'll change to reuse the flag definition in Arith dialect.

[Lewuathe]

Reuse the fastmath flag definitions in arith dialect.

[Lewuathe]

Remove unused header files.

[Lewuathe]

@kiranchandramohan @kuhar Could you give me your eyes on this change when you get a chance?

[kiranchandramohan]

I think we can reuse ArithFastMathInterface and ComplexFastMathInterface is not required. The functions in ComplexFastMathInterface seem to be the same as in ArithFastMathInterface. Was there any specific reason to define ComplexFastMathInterface? If not, I think reusing ArithFastMathInterface is the way forward.

mlir/include/mlir/Dialect/Complex/IR/CMakeLists.txt

9–12

I think the ComplexOpsInterface is not required.

mlir/include/mlir/Dialect/Complex/IR/Complex.h

28–33

This can be removed.

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

14

This can be replaced with ArithOpsInterfaces.td.

27

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

1

This Interface is not needed, I think. We can reuse ArithFastMathInterface.

mlir/lib/Conversion/ComplexToLLVM/ComplexToLLVM.cpp

60

This function seems to be available from mlir/Conversion/ArithCommon/AttrToLLVMConverter.h. Could you reuse that function?

You might have to add MLIRArithAttrToLLVMConversion if you go this route.

[Lewuathe]

Reuse interface and conversion function.

[Lewuathe]

Remove unued headers.

[kiranchandramohan]

Please update the test flang/test/Lower/complex-operations.f90 to fix the CI. Some minor comments inline.

LG. Please wait a couple of days before you submit to give others a chance to review.

mlir/lib/Conversion/ComplexToLLVM/ComplexToLLVM.cpp

203–206

Nit: Expand this type.

234–237

Nit: Expand this type.

282–285

Nit: Expand this type.

322–325

Nit: Expand this type.

mlir/lib/Dialect/Complex/IR/CMakeLists.txt

18

Nit: This was in alphabetical order.

mlir/test/Conversion/ComplexToLLVM/convert-to-llvm.mlir

259

Is a test for complex.abs missing?

[Lewuathe]

Be aware of the default Flang fastmath flag. https://reviews.llvm.org/D138675

[Lewuathe]

Post review follow-up

[Lewuathe]

Remove unused code.

[Lewuathe]

Specify the fastmath flag explicitly in flang tests.

[Lewuathe]

@kiranchandramohan Thank you so much for the review! I'll keep this change for a few days accordingly.

mlir/include/mlir/Conversion/ComplexToLLVM/ComplexToLLVM.h

 //===- ComplexToLLVM.h - Utils to convert from the complex dialect --------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 #ifndef MLIR_CONVERSION_COMPLEXTOLLVM_COMPLEXTOLLVM_H_
 #define MLIR_CONVERSION_COMPLEXTOLLVM_COMPLEXTOLLVM_H_
 
 #include "mlir/Conversion/LLVMCommon/StructBuilder.h"
+#include "mlir/Dialect/Complex/IR/Complex.h"
 
 namespace mlir {
 class LLVMTypeConverter;
 class Pass;
 class RewritePatternSet;
 
 #define GEN_PASS_DECL_CONVERTCOMPLEXTOLLVMPASS
 #include "mlir/Conversion/Passes.h.inc"

mlir/include/mlir/Dialect/Complex/IR/CMakeLists.txt

 add_mlir_dialect(ComplexOps complex)
 add_mlir_doc(ComplexOps ComplexOps Dialects/ -gen-dialect-doc)
 
 set(LLVM_TARGET_DEFINITIONS ComplexAttributes.td)
-mlir_tablegen(ComplexAttributes.h.inc -gen-attrdef-decls)
-mlir_tablegen(ComplexAttributes.cpp.inc -gen-attrdef-defs)
+mlir_tablegen(ComplexAttributes.h.inc -gen-attrdef-decls -attrdefs-dialect=complex)
+mlir_tablegen(ComplexAttributes.cpp.inc -gen-attrdef-defs -attrdefs-dialect=complex)
 add_public_tablegen_target(MLIRComplexAttributesIncGen)
+
+set(LLVM_TARGET_DEFINITIONS ComplexOpsInterfaces.td)
+mlir_tablegen(ComplexOpsInterfaces.h.inc -gen-op-interface-decls)
+mlir_tablegen(ComplexOpsInterfaces.cpp.inc -gen-op-interface-defs)
+add_public_tablegen_target(MLIRComplexOpsInterfacesIncGen)

mlir/include/mlir/Dialect/Complex/IR/Complex.h

 //===- Complex.h - Complex dialect --------------------------------*- C++-*-==//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_COMPLEX_IR_COMPLEX_H_
 #define MLIR_DIALECT_COMPLEX_IR_COMPLEX_H_
 
 #include "mlir/Bytecode/BytecodeOpInterface.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 
 //===----------------------------------------------------------------------===//
 // Complex Dialect
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Complex/IR/ComplexOpsDialect.h.inc"
 
+#define GET_ATTRDEF_CLASSES
+#include "mlir/Dialect/Complex/IR/ComplexAttributes.h.inc"
+
+//===----------------------------------------------------------------------===//
+// Complex Interfaces
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Complex/IR/ComplexOpsInterfaces.h.inc"
+
 //===----------------------------------------------------------------------===//
 // Complex Dialect Operations
 //===----------------------------------------------------------------------===//
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/Complex/IR/ComplexOps.h.inc"
 
-#define GET_ATTRDEF_CLASSES
-#include "mlir/Dialect/Complex/IR/ComplexAttributes.h.inc"
-
 #endif // MLIR_DIALECT_COMPLEX_IR_COMPLEX_H_

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

 //===- ComplexOps.td - Complex op definitions ----------------*- tablegen -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef COMPLEX_OPS
 #define COMPLEX_OPS
 
+include "mlir/Dialect/Arith/IR/ArithBase.td"
 include "mlir/Dialect/Complex/IR/ComplexBase.td"
+include "mlir/Dialect/Complex/IR/ComplexOpsInterfaces.td"
 include "mlir/IR/OpAsmInterface.td"
 include "mlir/Interfaces/InferTypeOpInterface.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 
 class Complex_Op<string mnemonic, list<Trait> traits = []>
     : Op<Complex_Dialect, mnemonic, traits>;
 
 // Base class for standard arithmetic operations on complex numbers with a
 // floating-point element type. These operations take two operands and return
 // one result, all of which must be complex numbers of the same type.
 class ComplexArithmeticOp<string mnemonic, list<Trait> traits = []> :
     Complex_Op<mnemonic, traits # [Pure, SameOperandsAndResultType,
-    Elementwise]> {
-  let arguments = (ins Complex<AnyFloat>:$lhs, Complex<AnyFloat>:$rhs);
+    Elementwise, DeclareOpInterfaceMethods<ComplexFastMathInterface>]> {
+  let arguments = (ins Complex<AnyFloat>:$lhs, Complex<AnyFloat>:$rhs, DefaultValuedAttr<
+        Arith_FastMathAttr, "::mlir::arith::FastMathFlags::none">:$fastmath);
   let results = (outs Complex<AnyFloat>:$result);
-  let assemblyFormat = "$lhs `,` $rhs  attr-dict `:` type($result)";
+  let assemblyFormat = "$lhs `,` $rhs (`fastmath` `` $fastmath^)? attr-dict `:` type($result)";
 }
 
 // Base class for standard unary operations on complex numbers with a
 // floating-point element type. These operations take one operand and return
 // one result; the operand must be a complex number.
 class ComplexUnaryOp<string mnemonic, list<Trait> traits = []> :
     Complex_Op<mnemonic, traits # [Pure, Elementwise]> {
-  let arguments = (ins Complex<AnyFloat>:$complex);
-  let assemblyFormat = "$complex attr-dict `:` type($complex)";
+  let arguments = (ins Complex<AnyFloat>:$complex, DefaultValuedAttr<
+        Arith_FastMathAttr, "::mlir::arith::FastMathFlags::none">:$fastmath);
+  let assemblyFormat = "$complex (`fastmath` `` $fastmath^)? attr-dict `:` type($complex)";
 }
 
 //===----------------------------------------------------------------------===//
 // AbsOp
 //===----------------------------------------------------------------------===//
 
 def AbsOp : ComplexUnaryOp<"abs",
     [TypesMatchWith<"complex element type matches result type",

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

+//===-- ComplexOpsInterfaces.td - complex op interfaces ---*- tablegen -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This is the interfaces definition file for complex dialect.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef COMPLEX_OPS_INTERFACES
+#define COMPLEX_OPS_INTERFACES
+
+include "mlir/IR/OpBase.td"
+include "mlir/Dialect/Complex/IR/ComplexBase.td"
+
+def ComplexFastMathInterface : OpInterface<"ComplexFastMathInterface"> {
+  let description = [{
+    Access to operation fastmath flags.
+  }];
+
+  let cppNamespace = "::mlir::complex";
+
+  let methods = [
+    InterfaceMethod<
+      /*desc=*/        "Returns a FastMathFlagsAttr attribute for the operation",
+      /*returnType=*/  "::mlir::arith::FastMathFlagsAttr",
+      /*methodName=*/  "getFastMathFlagsAttr",
+      /*args=*/        (ins),
+      /*methodBody=*/  [{}],
+      /*defaultImpl=*/ [{
+        ConcreteOp op = cast<ConcreteOp>(this->getOperation());
+        return op.getFastmathAttr();
+      }]
+      >,
+    StaticInterfaceMethod<
+      /*desc=*/        [{Returns the name of the FastMathFlagsAttr attribute
+                         for the operation}],
+      /*returnType=*/  "StringRef",
+      /*methodName=*/  "getFastMathAttrName",
+      /*args=*/        (ins),
+      /*methodBody=*/  [{}],
+      /*defaultImpl=*/ [{
+        return "fastmath";
+      }]
+      >
+
+  ];
+}
+
+#endif // COMPLEX_OPS_INTERFACES

mlir/lib/Conversion/ComplexToLLVM/ComplexToLLVM.cpp

 
 namespace mlir {
 #define GEN_PASS_DEF_CONVERTCOMPLEXTOLLVMPASS
 #include "mlir/Conversion/Passes.h.inc"
 } // namespace mlir
 
 using namespace mlir;
 using namespace mlir::LLVM;
+using namespace mlir::arith;
 
 //===----------------------------------------------------------------------===//
 // ComplexStructBuilder implementation.
 //===----------------------------------------------------------------------===//
 
 static constexpr unsigned kRealPosInComplexNumberStruct = 0;
 static constexpr unsigned kImaginaryPosInComplexNumberStruct = 1;
 
void ComplexStructBuilder::setImaginary(OpBuilder &builder, Location loc,
                                         Value imaginary) {
   setPtr(builder, loc, kImaginaryPosInComplexNumberStruct, imaginary);
 }
 
 Value ComplexStructBuilder::imaginary(OpBuilder &builder, Location loc) {
   return extractPtr(builder, loc, kImaginaryPosInComplexNumberStruct);
 }
 
+// Map arithmetic fastmath enum values to LLVMIR enum values.
+static LLVM::FastmathFlags
+convertComplexFastMathFlagsToLLVM(arith::FastMathFlags complexFMF) {
+  LLVM::FastmathFlags llvmFMF{};
+  const std::pair<arith::FastMathFlags, LLVM::FastmathFlags> flags[] = {
+      {arith::FastMathFlags::nnan, LLVM::FastmathFlags::nnan},
+      {arith::FastMathFlags::ninf, LLVM::FastmathFlags::ninf},
+      {arith::FastMathFlags::nsz, LLVM::FastmathFlags::nsz},
+      {arith::FastMathFlags::arcp, LLVM::FastmathFlags::arcp},
+      {arith::FastMathFlags::contract, LLVM::FastmathFlags::contract},
+      {arith::FastMathFlags::afn, LLVM::FastmathFlags::afn},
+      {arith::FastMathFlags::reassoc, LLVM::FastmathFlags::reassoc}};
+  for (auto fmfMap : flags) {
+    if (bitEnumContainsAny(complexFMF, fmfMap.first))
+      llvmFMF = llvmFMF | fmfMap.second;
+  }
+  return llvmFMF;
+}
+
 //===----------------------------------------------------------------------===//
 // Conversion patterns.
 //===----------------------------------------------------------------------===//
 
 namespace {
 
 struct AbsOpConversion : public ConvertOpToLLVMPattern<complex::AbsOp> {
   using ConvertOpToLLVMPattern<complex::AbsOp>::ConvertOpToLLVMPattern;
matchAndRewrite(complex::AddOp op, OpAdaptor adaptor,
     BinaryComplexOperands arg =
         unpackBinaryComplexOperands<complex::AddOp>(op, adaptor, rewriter);
 
     // Initialize complex number struct for result.
     auto structType = typeConverter->convertType(op.getType());
     auto result = ComplexStructBuilder::undef(rewriter, loc, structType);
 
     // Emit IR to add complex numbers.
-    auto fmf = LLVM::FastmathFlagsAttr::get(op.getContext(), {});
+    auto complexFMFAttr = op.getFastMathFlagsAttr();
+    auto fmf = LLVM::FastmathFlagsAttr::get(
+        op.getContext(),
+        convertComplexFastMathFlagsToLLVM(complexFMFAttr.getValue()));
     Value real =
         rewriter.create<LLVM::FAddOp>(loc, arg.lhs.real(), arg.rhs.real(), fmf);
     Value imag =
         rewriter.create<LLVM::FAddOp>(loc, arg.lhs.imag(), arg.rhs.imag(), fmf);
     result.setReal(rewriter, loc, real);
     result.setImaginary(rewriter, loc, imag);
 
     rewriter.replaceOp(op, {result});
matchAndRewrite(complex::DivOp op, OpAdaptor adaptor,
     BinaryComplexOperands arg =
         unpackBinaryComplexOperands<complex::DivOp>(op, adaptor, rewriter);
 
     // Initialize complex number struct for result.
     auto structType = typeConverter->convertType(op.getType());
     auto result = ComplexStructBuilder::undef(rewriter, loc, structType);
 
     // Emit IR to add complex numbers.
-    auto fmf = LLVM::FastmathFlagsAttr::get(op.getContext(), {});
+    auto complexFMFAttr = op.getFastMathFlagsAttr();
+    auto fmf = LLVM::FastmathFlagsAttr::get(
+        op.getContext(),
+        convertComplexFastMathFlagsToLLVM(complexFMFAttr.getValue()));
     Value rhsRe = arg.rhs.real();
     Value rhsIm = arg.rhs.imag();
     Value lhsRe = arg.lhs.real();
     Value lhsIm = arg.lhs.imag();
 
     Value rhsSqNorm = rewriter.create<LLVM::FAddOp>(
         loc, rewriter.create<LLVM::FMulOp>(loc, rhsRe, rhsRe, fmf),
         rewriter.create<LLVM::FMulOp>(loc, rhsIm, rhsIm, fmf), fmf);
matchAndRewrite(complex::MulOp op, OpAdaptor adaptor,
     BinaryComplexOperands arg =
         unpackBinaryComplexOperands<complex::MulOp>(op, adaptor, rewriter);
 
     // Initialize complex number struct for result.
     auto structType = typeConverter->convertType(op.getType());
     auto result = ComplexStructBuilder::undef(rewriter, loc, structType);
 
     // Emit IR to add complex numbers.
-    auto fmf = LLVM::FastmathFlagsAttr::get(op.getContext(), {});
+    auto complexFMFAttr = op.getFastMathFlagsAttr();
+    auto fmf = LLVM::FastmathFlagsAttr::get(
+        op.getContext(),
+        convertComplexFastMathFlagsToLLVM(complexFMFAttr.getValue()));
     Value rhsRe = arg.rhs.real();
     Value rhsIm = arg.rhs.imag();
     Value lhsRe = arg.lhs.real();
     Value lhsIm = arg.lhs.imag();
 
     Value real = rewriter.create<LLVM::FSubOp>(
         loc, rewriter.create<LLVM::FMulOp>(loc, rhsRe, lhsRe, fmf),
         rewriter.create<LLVM::FMulOp>(loc, rhsIm, lhsIm, fmf), fmf);
matchAndRewrite(complex::SubOp op, OpAdaptor adaptor,
     BinaryComplexOperands arg =
         unpackBinaryComplexOperands<complex::SubOp>(op, adaptor, rewriter);
 
     // Initialize complex number struct for result.
     auto structType = typeConverter->convertType(op.getType());
     auto result = ComplexStructBuilder::undef(rewriter, loc, structType);
 
     // Emit IR to substract complex numbers.
-    auto fmf = LLVM::FastmathFlagsAttr::get(op.getContext(), {});
+    auto complexFMFAttr = op.getFastMathFlagsAttr();
+    auto fmf = LLVM::FastmathFlagsAttr::get(
+        op.getContext(),
+        convertComplexFastMathFlagsToLLVM(complexFMFAttr.getValue()));
     Value real =
         rewriter.create<LLVM::FSubOp>(loc, arg.lhs.real(), arg.rhs.real(), fmf);
     Value imag =
         rewriter.create<LLVM::FSubOp>(loc, arg.lhs.imag(), arg.rhs.imag(), fmf);
     result.setReal(rewriter, loc, real);
     result.setImaginary(rewriter, loc, imag);
 
     rewriter.replaceOp(op, {result});

mlir/lib/Dialect/Complex/IR/CMakeLists.txt

 add_mlir_dialect_library(MLIRComplexDialect
   ComplexOps.cpp
   ComplexDialect.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Complex
 
   DEPENDS
   MLIRComplexOpsIncGen
   MLIRComplexAttributesIncGen
+  MLIRComplexOpsInterfacesIncGen
 
   LINK_LIBS PUBLIC
   MLIRArithDialect
   MLIRDialect
   MLIRInferTypeOpInterface
   MLIRIR
   )

mlir/test/Conversion/ComplexToLLVM/convert-to-llvm.mlir

 // CHECK: %[[REAL:.*]] = llvm.extractvalue %[[CASTED_ARG]][0] : ![[C_TY]]
 // CHECK: %[[IMAG:.*]] = llvm.extractvalue %[[CASTED_ARG]][1] : ![[C_TY]]
 // CHECK-DAG: %[[REAL_SQ:.*]] = llvm.fmul %[[REAL]], %[[REAL]]  : f32
 // CHECK-DAG: %[[IMAG_SQ:.*]] = llvm.fmul %[[IMAG]], %[[IMAG]]  : f32
 // CHECK: %[[SQ_NORM:.*]] = llvm.fadd %[[REAL_SQ]], %[[IMAG_SQ]]  : f32
 // CHECK: %[[NORM:.*]] = llvm.intr.sqrt(%[[SQ_NORM]]) : (f32) -> f32
 // CHECK: return %[[NORM]] : f32
 
+// CHECK-LABEL: func @complex_addition_with_fmf
+// CHECK-DAG:     %[[A_REAL:.*]] = llvm.extractvalue %[[A:.*]][0] : !llvm.struct<(f64, f64)>
+// CHECK-DAG:     %[[B_REAL:.*]] = llvm.extractvalue %[[B:.*]][0] : !llvm.struct<(f64, f64)>
+// CHECK-DAG:     %[[A_IMAG:.*]] = llvm.extractvalue %[[A]][1] : !llvm.struct<(f64, f64)>
+// CHECK-DAG:     %[[B_IMAG:.*]] = llvm.extractvalue %[[B]][1] : !llvm.struct<(f64, f64)>
+// CHECK:         %[[C0:.*]] = llvm.mlir.undef : !llvm.struct<(f64, f64)>
+// CHECK-DAG:     %[[C_REAL:.*]] = llvm.fadd %[[A_REAL]], %[[B_REAL]] {fastmathFlags = #llvm.fastmath<reassoc>} : f64
+// CHECK-DAG:     %[[C_IMAG:.*]] = llvm.fadd %[[A_IMAG]], %[[B_IMAG]] {fastmathFlags = #llvm.fastmath<reassoc>} : f64
+// CHECK:         %[[C1:.*]] = llvm.insertvalue %[[C_REAL]], %[[C0]][0] : !llvm.struct<(f64, f64)>
+// CHECK:         %[[C2:.*]] = llvm.insertvalue %[[C_IMAG]], %[[C1]][1] : !llvm.struct<(f64, f64)>
+func.func @complex_addition_with_fmf() {
+  %a_re = arith.constant 1.2 : f64
+  %a_im = arith.constant 3.4 : f64
+  %a = complex.create %a_re, %a_im : complex<f64>
+  %b_re = arith.constant 5.6 : f64
+  %b_im = arith.constant 7.8 : f64
+  %b = complex.create %b_re, %b_im : complex<f64>
+  %c = complex.add %a, %b fastmath<reassoc> : complex<f64>
+  return
+}
+
+// CHECK-LABEL: func @complex_substraction_with_fmf
+// CHECK-DAG:     %[[A_REAL:.*]] = llvm.extractvalue %[[A:.*]][0] : !llvm.struct<(f64, f64)>
+// CHECK-DAG:     %[[B_REAL:.*]] = llvm.extractvalue %[[B:.*]][0] : !llvm.struct<(f64, f64)>
+// CHECK-DAG:     %[[A_IMAG:.*]] = llvm.extractvalue %[[A]][1] : !llvm.struct<(f64, f64)>
+// CHECK-DAG:     %[[B_IMAG:.*]] = llvm.extractvalue %[[B]][1] : !llvm.struct<(f64, f64)>
+// CHECK:         %[[C0:.*]] = llvm.mlir.undef : !llvm.struct<(f64, f64)>
+// CHECK-DAG:     %[[C_REAL:.*]] = llvm.fsub %[[A_REAL]], %[[B_REAL]] {fastmathFlags = #llvm.fastmath<nnan, ninf>} : f64
+// CHECK-DAG:     %[[C_IMAG:.*]] = llvm.fsub %[[A_IMAG]], %[[B_IMAG]] {fastmathFlags = #llvm.fastmath<nnan, ninf>} : f64
+// CHECK:         %[[C1:.*]] = llvm.insertvalue %[[C_REAL]], %[[C0]][0] : !llvm.struct<(f64, f64)>
+// CHECK:         %[[C2:.*]] = llvm.insertvalue %[[C_IMAG]], %[[C1]][1] : !llvm.struct<(f64, f64)>
+func.func @complex_substraction_with_fmf() {
+  %a_re = arith.constant 1.2 : f64
+  %a_im = arith.constant 3.4 : f64
+  %a = complex.create %a_re, %a_im : complex<f64>
+  %b_re = arith.constant 5.6 : f64
+  %b_im = arith.constant 7.8 : f64
+  %b = complex.create %b_re, %b_im : complex<f64>
+  %c = complex.sub %a, %b fastmath<nnan,ninf> : complex<f64>
+  return
+}
+
+// CHECK-LABEL: func @complex_div_with_fmf
+// CHECK-SAME:    %[[LHS:.*]]: complex<f32>, %[[RHS:.*]]: complex<f32>
+// CHECK-DAG: %[[CASTED_LHS:.*]] = builtin.unrealized_conversion_cast %[[LHS]] : complex<f32> to ![[C_TY:.*>]]
+// CHECK-DAG: %[[CASTED_RHS:.*]] = builtin.unrealized_conversion_cast %[[RHS]] : complex<f32> to ![[C_TY]]
+
+// CHECK: %[[LHS_RE:.*]] = llvm.extractvalue %[[CASTED_LHS]][0] : ![[C_TY]]
+// CHECK: %[[LHS_IM:.*]] = llvm.extractvalue %[[CASTED_LHS]][1] : ![[C_TY]]
+// CHECK: %[[RHS_RE:.*]] = llvm.extractvalue %[[CASTED_RHS]][0] : ![[C_TY]]
+// CHECK: %[[RHS_IM:.*]] = llvm.extractvalue %[[CASTED_RHS]][1] : ![[C_TY]]
+
+// CHECK: %[[RESULT_0:.*]] = llvm.mlir.undef : ![[C_TY]]
+
+// CHECK-DAG: %[[RHS_RE_SQ:.*]] = llvm.fmul %[[RHS_RE]], %[[RHS_RE]] {fastmathFlags = #llvm.fastmath<nsz, arcp>} : f32
+// CHECK-DAG: %[[RHS_IM_SQ:.*]] = llvm.fmul %[[RHS_IM]], %[[RHS_IM]] {fastmathFlags = #llvm.fastmath<nsz, arcp>} : f32
+// CHECK: %[[SQ_NORM:.*]] = llvm.fadd %[[RHS_RE_SQ]], %[[RHS_IM_SQ]] {fastmathFlags = #llvm.fastmath<nsz, arcp>} : f32
+
+// CHECK-DAG: %[[REAL_TMP_0:.*]] = llvm.fmul %[[LHS_RE]], %[[RHS_RE]] {fastmathFlags = #llvm.fastmath<nsz, arcp>} : f32
+// CHECK-DAG: %[[REAL_TMP_1:.*]] = llvm.fmul %[[LHS_IM]], %[[RHS_IM]] {fastmathFlags = #llvm.fastmath<nsz, arcp>} : f32
+// CHECK: %[[REAL_TMP_2:.*]] = llvm.fadd %[[REAL_TMP_0]], %[[REAL_TMP_1]] {fastmathFlags = #llvm.fastmath<nsz, arcp>} : f32
+
+// CHECK-DAG: %[[IMAG_TMP_0:.*]] = llvm.fmul %[[LHS_IM]], %[[RHS_RE]] {fastmathFlags = #llvm.fastmath<nsz, arcp>} : f32
+// CHECK-DAG: %[[IMAG_TMP_1:.*]] = llvm.fmul %[[LHS_RE]], %[[RHS_IM]] {fastmathFlags = #llvm.fastmath<nsz, arcp>} : f32
+// CHECK: %[[IMAG_TMP_2:.*]] = llvm.fsub %[[IMAG_TMP_0]], %[[IMAG_TMP_1]] {fastmathFlags = #llvm.fastmath<nsz, arcp>} : f32
+
+// CHECK: %[[REAL:.*]] = llvm.fdiv %[[REAL_TMP_2]], %[[SQ_NORM]] {fastmathFlags = #llvm.fastmath<nsz, arcp>} : f32
+// CHECK: %[[RESULT_1:.*]] = llvm.insertvalue %[[REAL]], %[[RESULT_0]][0] : ![[C_TY]]
+// CHECK: %[[IMAG:.*]] = llvm.fdiv %[[IMAG_TMP_2]], %[[SQ_NORM]] {fastmathFlags = #llvm.fastmath<nsz, arcp>} : f32
+// CHECK: %[[RESULT_2:.*]] = llvm.insertvalue %[[IMAG]], %[[RESULT_1]][1] : ![[C_TY]]
+//
+// CHECK: %[[CASTED_RESULT:.*]] = builtin.unrealized_conversion_cast %[[RESULT_2]] : ![[C_TY]] to complex<f32>
+// CHECK: return %[[CASTED_RESULT]] : complex<f32>
+func.func @complex_div_with_fmf(%lhs: complex<f32>, %rhs: complex<f32>) -> complex<f32> {
+  %div = complex.div %lhs, %rhs fastmath<nsz,arcp> : complex<f32>
+  return %div : complex<f32>
+}
+
+
+// CHECK-LABEL: func @complex_mul_with_fmf
+// CHECK-SAME:    %[[LHS:.*]]: complex<f32>, %[[RHS:.*]]: complex<f32>
+// CHECK-DAG: %[[CASTED_LHS:.*]] = builtin.unrealized_conversion_cast %[[LHS]] : complex<f32> to ![[C_TY:.*>]]
+// CHECK-DAG: %[[CASTED_RHS:.*]] = builtin.unrealized_conversion_cast %[[RHS]] : complex<f32> to ![[C_TY]]
+
+// CHECK: %[[LHS_RE:.*]] = llvm.extractvalue %[[CASTED_LHS]][0] : ![[C_TY]]
+// CHECK: %[[LHS_IM:.*]] = llvm.extractvalue %[[CASTED_LHS]][1] : ![[C_TY]]
+// CHECK: %[[RHS_RE:.*]] = llvm.extractvalue %[[CASTED_RHS]][0] : ![[C_TY]]
+// CHECK: %[[RHS_IM:.*]] = llvm.extractvalue %[[CASTED_RHS]][1] : ![[C_TY]]
+// CHECK: %[[RESULT_0:.*]] = llvm.mlir.undef : ![[C_TY]]
+
+// CHECK-DAG: %[[REAL_TMP_0:.*]] = llvm.fmul %[[RHS_RE]], %[[LHS_RE]] {fastmathFlags = #llvm.fastmath<contract, afn>} : f32
+// CHECK-DAG: %[[REAL_TMP_1:.*]] = llvm.fmul %[[RHS_IM]], %[[LHS_IM]] {fastmathFlags = #llvm.fastmath<contract, afn>} : f32
+// CHECK: %[[REAL:.*]] = llvm.fsub %[[REAL_TMP_0]], %[[REAL_TMP_1]] {fastmathFlags = #llvm.fastmath<contract, afn>} : f32
+
+// CHECK-DAG: %[[IMAG_TMP_0:.*]] = llvm.fmul %[[LHS_IM]], %[[RHS_RE]] {fastmathFlags = #llvm.fastmath<contract, afn>} : f32
+// CHECK-DAG: %[[IMAG_TMP_1:.*]] = llvm.fmul %[[LHS_RE]], %[[RHS_IM]] {fastmathFlags = #llvm.fastmath<contract, afn>} : f32
+// CHECK: %[[IMAG:.*]] = llvm.fadd %[[IMAG_TMP_0]], %[[IMAG_TMP_1]] {fastmathFlags = #llvm.fastmath<contract, afn>} : f32
+
+// CHECK: %[[RESULT_1:.*]] = llvm.insertvalue %[[REAL]], %[[RESULT_0]][0]
+// CHECK: %[[RESULT_2:.*]] = llvm.insertvalue %[[IMAG]], %[[RESULT_1]][1]
+
+// CHECK: %[[CASTED_RESULT:.*]] = builtin.unrealized_conversion_cast %[[RESULT_2]] : ![[C_TY]] to complex<f32>
+// CHECK: return %[[CASTED_RESULT]] : complex<f32>
+func.func @complex_mul_with_fmf(%lhs: complex<f32>, %rhs: complex<f32>) -> complex<f32> {
+  %mul = complex.mul %lhs, %rhs fastmath<contract,afn>  : complex<f32>
+  return %mul : complex<f32>
+}