[mlir][complex] Canonicalize complex.div by one (#85513)

We can canonicalize the complex.div if the divisor is one (real = 1.0,
imag = 0.0) with the input number itself.

Ref: https://www.cuemath.com/numbers/division-of-complex-numbers/

1 files changed, 26 insertions, 0 deletions

diff --git a/mlir/lib/Dialect/Complex/IR/ComplexOps.cpp b/mlir/lib/Dialect/Complex/IR/ComplexOps.cpp
index 5529dcc..1c81433 100644
--- a/mlir/lib/Dialect/Complex/IR/ComplexOps.cpp
+++ b/mlir/lib/Dialect/Complex/IR/ComplexOps.cpp
@@ -370,6 +370,32 @@ OpFoldResult MulOp::fold(FoldAdaptor adaptor) {
 }
 
 //===----------------------------------------------------------------------===//
+// DivOp
+//===----------------------------------------------------------------------===//
+
+OpFoldResult DivOp::fold(FoldAdaptor adaptor) {
+  auto rhs = adaptor.getRhs();
+  if (!rhs)
+    return {};
+
+  ArrayAttr arrayAttr = rhs.dyn_cast<ArrayAttr>();
+  if (!arrayAttr || arrayAttr.size() != 2)
+    return {};
+
+  APFloat real = cast<FloatAttr>(arrayAttr[0]).getValue();
+  APFloat imag = cast<FloatAttr>(arrayAttr[1]).getValue();
+
+  if (!imag.isZero())
+    return {};
+
+  // complex.div(a, complex.constant<1.0, 0.0>) -> a
+  if (real == APFloat(real.getSemantics(), 1))
+    return getLhs();
+
+  return {};
+}
+
+//===----------------------------------------------------------------------===//
 // TableGen'd op method definitions
 //===----------------------------------------------------------------------===//