[mlir][complex] Canonicalize complex.mul with 1 and 0

[Lewuathe]

Summary

We can fold the complex.mul if the right value is obvious 1 or 0.

[kuhar]

Does complex allow for special fp values like NaN, +/- Inf, or -0? I remember this being the case for complex in C. If yes, is the '* 0' fold valid when the non-constant values are NaN/Inf/-0/etc.?

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

256

Prefer early exits: https://llvm.org/docs/CodingStandards.html#use-early-exits-and-continue-to-simplify-code.

Also in the code below.

257–259

nit: I would prefer to see the exact types spelled out here, it's not immediately obvious what each .getValue(); returns.

[jfurtek]

As @kuhar mentioned, multiplying by zero isn't necessarily zero, unless the appropriate fastmath flags are set. (LLVM won't simplify multiplication scalar values by zero unless the nnan and nsz flags are set.)

I would think the "right" thing to do here is add fastmath flags to the complex dialect operations and make the fold conditional.

[Lewuathe]

@kuhar @jfurtek Thank you for the comment. I found a related discussion here for the arithmetic dialect. (you are already leading the one, @jfurtek)

https://discourse.llvm.org/t/rfc-fastmath-flags-support-in-mlir-arith-dialect/6049

I think we can add similar flag as the arith dialect right?

IIRC mul_one case should not be affected by the fast math flag this time. Can I separate mul_zero from the patch so that we can work on after supporting fast math flag?

[kuhar]

In D153606#4445249, @Lewuathe wrote:

@kuhar @jfurtek Thank you for the comment. I found a related discussion here for the arithmetic dialect. (you are already leading the one, @jfurtek)

IIRC mul_one case should not be affected by the fast math flag this time. Can I separate mul_zero from the patch so that we can work on after supporting fast math flag?

Limiting the scope SGTM. We should make sure to add test cases with those special fp values to show that these folds handle them as intended, provided they won't be matched against directly in the fold.

[Lewuathe]

Add test cases for multiple floating point precisions.

[kuhar]

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

261–262

nit: we don't need llvm:: with casts, I believe it started appearing in the code base because it was easier to migrate from member casts this way

264–265

Consider adding an early exit if (real != APFloat(...)) return {}; to simplify the switch below

[Lewuathe]

Post review follow-up

[kuhar]

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

257–259

The type of arrayAttr is not obvious to me based on the RHS. Could we spell it out instead of using auto?

264–265

Not done

[Lewuathe]

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

264–265

@kuhar Sorry, I might have missed what you meant. Which part should we do early exit for this case?

To do if (real != APFloat(...)) type of check, we need to check the size in bits first (64 bit real will always fail to check real != APFloat(1.0f) but it should be caught in the case of 32 bit. )

[kuhar]

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

264–265

I was thinking of something like:

if (!imag.isZero()) return {};
if (real != APFloat(1.0)) return {};

if (!llvm::is_contained({32, 64}, real.getSizeInBits(real.getSemantics())) return {};

return getLhs();

[kuhar]

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

264–265

Gah, I'm sorry, I missed the 1.0 vs. 1.0f distinction.

In this case I think we can use the m_OneFloat matcher: https://mlir.llvm.org/doxygen/namespacemlir.html#af0495d84f34cf3238a7741fa6974a485

[Lewuathe]

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

264–265

Thank you. Let me try!

[Lewuathe]

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

264–265

Apparently, m_OneFloat (or any other matcher) does not work with APFloat derived from the attribute. It's only working with the value type.

https://discord.com/channels/636084430946959380/642426447167881246/1123498353129426994

[Lewuathe]

Early exit for non-zero imaginary part.

[Lewuathe]

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

264–265

I changed to exit earlier in the case of non-zero imaginary part instead.

[kuhar]

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

264–265

In this case, should we use the same implementation as they do to support any bitwidth?

inline detail::constant_float_predicate_matcher m_OneFloat() {
  return {[](const APFloat &value) {
    return APFloat(value.getSemantics(), 1) == value;
  }};
}

[Lewuathe]

Support any bitwidth (f16, f80 and f128)

[kuhar]

LGTM, thanks for the fixes!

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

def MulOp : ComplexArithmeticOp<"mul"> {
   let summary = "complex multiplication";
   let description = [{
     The `mul` operation takes two complex numbers and returns their product:
 
     ```mlir
     %a = complex.mul %b, %c : complex<f32>
     ```
   }];
+
+  let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // NegOp
 //===----------------------------------------------------------------------===//
 
 def NegOp : ComplexUnaryOp<"neg", [SameOperandsAndResultType]> {
   let summary = "Negation operator";

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

OpFoldResult ConjOp::fold(FoldAdaptor adaptor) {
   // complex.conj(complex.conj(a)) -> a
   if (auto conjOp = getOperand().getDefiningOp<ConjOp>())
     return conjOp.getOperand();
 
   return {};
 }
 
 //===----------------------------------------------------------------------===//
+// MulOp
+//===----------------------------------------------------------------------===//
+
+OpFoldResult MulOp::fold(FoldAdaptor adaptor) {
+  if (auto constant = getRhs().getDefiningOp<ConstantOp>()) {
+    auto arrayAttr = constant.getValue();
+    auto real = llvm::cast<FloatAttr>(arrayAttr[0]).getValue();
+    auto imag = llvm::cast<FloatAttr>(arrayAttr[1]).getValue();
+
+    if (imag.isZero()) {
+      // complex.mul(a, complex.constant<0.0, 0.0>)
+      //   -> complex.constant<0.0, 0.0>
+      if (real.isZero())
+        return getRhs();
+
+      // complex.mul(a, complex.constant<1.0, 0.0>) -> a
+      switch (real.getSizeInBits(real.getSemantics())) {
+      case 64:
+        if (real == APFloat(1.0))
+          return getLhs();
+        return {};
+      case 32:
+        if (real == APFloat(1.0f))
+          return getLhs();
+        return {};
+      default:
+        return {};
+      }
+    }
+  }
+
+  return {};
+}
+
+//===----------------------------------------------------------------------===//
 // TableGen'd op method definitions
 //===----------------------------------------------------------------------===//
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/Complex/IR/ComplexOps.cpp.inc"

mlir/test/Dialect/Complex/canonicalize.mlir

 func.func @im_neg(%arg0: f32, %arg1: f32) -> f32 {
   %create = complex.create %arg0, %arg1: complex<f32>
   // CHECK: %[[NEG:.*]] = arith.negf %[[ARG1]]
   %neg = complex.neg %create : complex<f32>
   %im = complex.im %neg : complex<f32>
   // CHECK-NEXT: return %[[NEG]]
   return %im : f32
 }
+
+// CHECK-LABEL: func @mul_zero
+//  CHECK-SAME: (%[[ARG0:.*]]: f32, %[[ARG1:.*]]: f32) -> complex<f32>
+func.func @mul_zero(%arg0: f32, %arg1: f32) -> complex<f32> {
+  %create = complex.create %arg0, %arg1: complex<f32>  
+  %zero = complex.constant [0.0 : f32, 0.0 : f32] : complex<f32>
+  %mul = complex.mul %create, %zero : complex<f32>
+  // CHECK: %[[ZERO:.*]] = complex.constant [0.000000e+00 : f32, 0.000000e+00 : f32] : complex<f32>
+  // CHECK-NEXT: return %[[ZERO]]
+  return %mul : complex<f32>
+}
+
+// CHECK-LABEL: func @mul_one_f32
+//  CHECK-SAME: (%[[ARG0:.*]]: f32, %[[ARG1:.*]]: f32) -> complex<f32>
+func.func @mul_one_f32(%arg0: f32, %arg1: f32) -> complex<f32> {
+  %create = complex.create %arg0, %arg1: complex<f32>  
+  %one = complex.constant [1.0 : f32, 0.0 : f32] : complex<f32>
+  %mul = complex.mul %create, %one : complex<f32>
+  // CHECK: %[[CREATE:.*]] = complex.create %[[ARG0]], %[[ARG1]] : complex<f32>
+  // CHECK-NEXT: return %[[CREATE]]
+  return %mul : complex<f32>
+}
+
+// CHECK-LABEL: func @mul_one_f64
+//  CHECK-SAME: (%[[ARG0:.*]]: f64, %[[ARG1:.*]]: f64) -> complex<f64>
+func.func @mul_one_f64(%arg0: f64, %arg1: f64) -> complex<f64> {
+  %create = complex.create %arg0, %arg1: complex<f64>  
+  %one = complex.constant [1.0 : f64, 0.0 : f64] : complex<f64>
+  %mul = complex.mul %create, %one : complex<f64>
+  // CHECK: %[[CREATE:.*]] = complex.create %[[ARG0]], %[[ARG1]] : complex<f64>
+  // CHECK-NEXT: return %[[CREATE]]
+  return %mul : complex<f64>
+}