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

[Kai Sasaki]

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

>From 3bcec8cd3df62ceec68bd0f68cb4741a9d41cd2b 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] 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
 }