[Mlir-commits] [mlir] [mlir][complex] Fastmath integer overflow flag in complex dialect ops (PR #90823)

[Kai Sasaki]

https://github.com/Lewuathe created https://github.com/llvm/llvm-project/pull/90823

We are going to support the integer overflow flag so that we can make use of this information for further optimization and canonicalization. Although the complex dialect does not use any `AddI`, `SubI` or `MulI` in the arith dialect diectly in the conversion to the standard dialect, we can utilize the flag in the conversion to LLVM or other dialect.

This PR adds the interface of the int overflow flag and changes the assembly format to allow the flag to be incoming. 

See: https://discourse.llvm.org/t/rfc-integer-overflow-flags-support-in-arith-dialect

>From 37dd9cd979425143f7824d55b8af0bafe37b6185 Mon Sep 17 00:00:00 2001
From: Kai Sasaki <lewuathe at gmail.com>
Date: Thu, 2 May 2024 15:01:19 +0900
Subject: [PATCH] [mlir][complex] Fastmath integer overflow flag in complex
 dialect ops

---
 .../mlir/Dialect/Complex/IR/ComplexOps.td     | 19 +++++++++++--------
 mlir/test/Dialect/Complex/ops.mlir            |  6 ++++++
 2 files changed, 17 insertions(+), 8 deletions(-)

diff --git a/mlir/include/mlir/Dialect/Complex/IR/ComplexOps.td b/mlir/include/mlir/Dialect/Complex/IR/ComplexOps.td
index d660292478b190..9b492d96a34253 100644
--- a/mlir/include/mlir/Dialect/Complex/IR/ComplexOps.td
+++ b/mlir/include/mlir/Dialect/Complex/IR/ComplexOps.td
@@ -24,21 +24,24 @@ class Complex_Op<string mnemonic, list<Trait> traits = []>
 // 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, DeclareOpInterfaceMethods<ArithFastMathInterface>]> {
-  let arguments = (ins Complex<AnyFloat>:$lhs, Complex<AnyFloat>:$rhs, DefaultValuedAttr<
-        Arith_FastMathAttr, "::mlir::arith::FastMathFlags::none">:$fastmath);
+    Elementwise, DeclareOpInterfaceMethods<ArithFastMathInterface>, DeclareOpInterfaceMethods<ArithIntegerOverflowFlagsInterface>]> {
+  let arguments = (ins Complex<AnyFloat>:$lhs, Complex<AnyFloat>:$rhs, 
+        DefaultValuedAttr<Arith_FastMathAttr, "::mlir::arith::FastMathFlags::none">:$fastmath,
+        DefaultValuedAttr<Arith_IntegerOverflowAttr,"::mlir::arith::IntegerOverflowFlags::none">:$overflowFlags);
   let results = (outs Complex<AnyFloat>:$result);
-  let assemblyFormat = "$lhs `,` $rhs (`fastmath` `` $fastmath^)? attr-dict `:` type($result)";
+  let assemblyFormat = "$lhs `,` $rhs (`fastmath` `` $fastmath^)? (`overflow` `` $overflowFlags^)? 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, DeclareOpInterfaceMethods<ArithFastMathInterface>]> {
-  let arguments = (ins Complex<AnyFloat>:$complex, DefaultValuedAttr<
-        Arith_FastMathAttr, "::mlir::arith::FastMathFlags::none">:$fastmath);
-  let assemblyFormat = "$complex (`fastmath` `` $fastmath^)? attr-dict `:` type($complex)";
+    Complex_Op<mnemonic, traits # [Pure, Elementwise, 
+      DeclareOpInterfaceMethods<ArithFastMathInterface>, DeclareOpInterfaceMethods<ArithIntegerOverflowFlagsInterface>]> {
+  let arguments = (ins Complex<AnyFloat>:$complex, 
+        DefaultValuedAttr<Arith_FastMathAttr, "::mlir::arith::FastMathFlags::none">:$fastmath,
+        DefaultValuedAttr<Arith_IntegerOverflowAttr,"::mlir::arith::IntegerOverflowFlags::none">:$overflowFlags);
+  let assemblyFormat = "$complex (`fastmath` `` $fastmath^)? (`overflow` `` $overflowFlags^)? attr-dict `:` type($complex)";
 }
 
 //===----------------------------------------------------------------------===//
diff --git a/mlir/test/Dialect/Complex/ops.mlir b/mlir/test/Dialect/Complex/ops.mlir
index 96f17b2898c834..597bd00feb801e 100644
--- a/mlir/test/Dialect/Complex/ops.mlir
+++ b/mlir/test/Dialect/Complex/ops.mlir
@@ -89,5 +89,11 @@ func.func @ops(%f: f32) {
   // CHECK: complex.bitcast %[[C]]
   %i64 = complex.bitcast %complex : complex<f32> to i64
 
+ // CHECK: complex.add %[[C]], %[[C]] overflow<nsw> : complex<f32>
+  %add_intflags = complex.add %complex, %complex overflow<nsw> : complex<f32>  
+
+ // CHECK: complex.sub %[[C]], %[[C]] overflow<nsw, nuw> : complex<f32>
+  %sub_intflags = complex.sub %complex, %complex overflow<nsw, nuw> : complex<f32>  
+
   return
 }