[Mlir-commits] [mlir] [mlir][complex] Canonicalize complex.div by one (PR #85513)

[Kai Sasaki]

https://github.com/Lewuathe created https://github.com/llvm/llvm-project/pull/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/

>From 64b4735afa67e31a782bc19cb748c9abf0164a2a Mon Sep 17 00:00:00 2001
From: Kai Sasaki <lewuathe at gmail.com>
Date: Tue, 12 Mar 2024 16:26:59 +0900
Subject: [PATCH] [mlir][complex] Canonicalize complex.div by one

---
 .../mlir/Dialect/Complex/IR/ComplexOps.td     |  2 +
 mlir/lib/Dialect/Complex/IR/ComplexOps.cpp    | 23 ++++++++
 mlir/test/Dialect/Complex/canonicalize.mlir   | 56 +++++++++++++++++++
 3 files changed, 81 insertions(+)

diff --git a/mlir/include/mlir/Dialect/Complex/IR/ComplexOps.td b/mlir/include/mlir/Dialect/Complex/IR/ComplexOps.td
index e19d714cadf8a9..d660292478b190 100644
--- a/mlir/include/mlir/Dialect/Complex/IR/ComplexOps.td
+++ b/mlir/include/mlir/Dialect/Complex/IR/ComplexOps.td
@@ -225,6 +225,8 @@ def DivOp : ComplexArithmeticOp<"div"> {
     %a = complex.div %b, %c : complex<f32>
     ```
   }];
+
+  let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Dialect/Complex/IR/ComplexOps.cpp b/mlir/lib/Dialect/Complex/IR/ComplexOps.cpp
index 5529dccaf1de6f..bb0ac17727d828 100644
--- a/mlir/lib/Dialect/Complex/IR/ComplexOps.cpp
+++ b/mlir/lib/Dialect/Complex/IR/ComplexOps.cpp
@@ -369,6 +369,29 @@ OpFoldResult MulOp::fold(FoldAdaptor adaptor) {
   return {};
 }
 
+//===----------------------------------------------------------------------===//
+// DivOp
+//===----------------------------------------------------------------------===//
+
+OpFoldResult DivOp::fold(FoldAdaptor adaptor) {
+  auto constant = getRhs().getDefiningOp<ConstantOp>();
+  if (!constant)
+    return {};
+
+  ArrayAttr arrayAttr = constant.getValue();
+  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
 //===----------------------------------------------------------------------===//
diff --git a/mlir/test/Dialect/Complex/canonicalize.mlir b/mlir/test/Dialect/Complex/canonicalize.mlir
index 64c3f313dda959..16fa0fdf56a1b4 100644
--- a/mlir/test/Dialect/Complex/canonicalize.mlir
+++ b/mlir/test/Dialect/Complex/canonicalize.mlir
@@ -281,3 +281,59 @@ func.func @double_reverse_bitcast(%arg0 : complex<f32>) -> f64 {
   // CHECK: return %[[R0]] : f64
   func.return %1 : f64
 }
+
+
+// CHECK-LABEL: func @div_one_f16
+//  CHECK-SAME: (%[[ARG0:.*]]: f16, %[[ARG1:.*]]: f16) -> complex<f16>
+func.func @div_one_f16(%arg0: f16, %arg1: f16) -> complex<f16> {
+  %create = complex.create %arg0, %arg1: complex<f16>  
+  %one = complex.constant [1.0 : f16, 0.0 : f16] : complex<f16>
+  %div = complex.div %create, %one : complex<f16>
+  // CHECK: %[[CREATE:.*]] = complex.create %[[ARG0]], %[[ARG1]] : complex<f16>
+  // CHECK-NEXT: return %[[CREATE]]
+  return %div : complex<f16>
+}
+
+// CHECK-LABEL: func @div_one_f32
+//  CHECK-SAME: (%[[ARG0:.*]]: f32, %[[ARG1:.*]]: f32) -> complex<f32>
+func.func @div_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>
+  %div = complex.div %create, %one : complex<f32>
+  // CHECK: %[[CREATE:.*]] = complex.create %[[ARG0]], %[[ARG1]] : complex<f32>
+  // CHECK-NEXT: return %[[CREATE]]
+  return %div : complex<f32>
+}
+
+// CHECK-LABEL: func @div_one_f64
+//  CHECK-SAME: (%[[ARG0:.*]]: f64, %[[ARG1:.*]]: f64) -> complex<f64>
+func.func @div_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>
+  %div = complex.div %create, %one : complex<f64>
+  // CHECK: %[[CREATE:.*]] = complex.create %[[ARG0]], %[[ARG1]] : complex<f64>
+  // CHECK-NEXT: return %[[CREATE]]
+  return %div : complex<f64>
+}
+
+// CHECK-LABEL: func @div_one_f80
+//  CHECK-SAME: (%[[ARG0:.*]]: f80, %[[ARG1:.*]]: f80) -> complex<f80>
+func.func @div_one_f80(%arg0: f80, %arg1: f80) -> complex<f80> {
+  %create = complex.create %arg0, %arg1: complex<f80>  
+  %one = complex.constant [1.0 : f80, 0.0 : f80] : complex<f80>
+  %div = complex.div %create, %one : complex<f80>
+  // CHECK: %[[CREATE:.*]] = complex.create %[[ARG0]], %[[ARG1]] : complex<f80>
+  // CHECK-NEXT: return %[[CREATE]]
+  return %div : complex<f80>
+}
+
+// CHECK-LABEL: func @div_one_f128
+//  CHECK-SAME: (%[[ARG0:.*]]: f128, %[[ARG1:.*]]: f128) -> complex<f128>
+func.func @div_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>
+  %div = complex.div %create, %one : complex<f128>
+  // CHECK: %[[CREATE:.*]] = complex.create %[[ARG0]], %[[ARG1]] : complex<f128>
+  // CHECK-NEXT: return %[[CREATE]]
+  return %div : complex<f128>
+}