[mlir][complex] Add Integer overflow flag in complex dialect ops

[Lewuathe]

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

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-complex

Author: Kai Sasaki (Lewuathe)

Changes

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

---

Full diff: https://github.com/llvm/llvm-project/pull/90823.diff

2 Files Affected:

• (modified) mlir/include/mlir/Dialect/Complex/IR/ComplexOps.td (+11-8)
• (modified) mlir/test/Dialect/Complex/ops.mlir (+6)

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
 }

[Lewuathe]

@joker-eph @matthias-springer Can I ask you to review this PR when you get a chance?

[joker-eph]

Can you add a lowering test as well?

[Lewuathe]

@joker-eph Thanks for taking a look at this. I found the integer overflow flag is only supported by arith.addi, arith.subi, arith.muli and arith.shli for now. The complex dialect does not use any these ops to lower.

Although I thought we have needed to support it or provided the interface for future use case, it might have been for this dialect?