From 6e951b3ec99e5cc9d72e532ead80de4af666d60d Mon Sep 17 00:00:00 2001
From: Adrian Kuegel <akuegel@google.com>
Date: Fri, 29 Jul 2022 13:30:29 +0200
Subject: [PATCH] [mlir][Complex] Add convenience builder for complex.number
 attribute.

Differential Revision: https://reviews.llvm.org/D130756
---
 .../Dialect/Complex/IR/ComplexAttributes.td     | 17 +++++++++++++++++
 mlir/lib/Dialect/Complex/IR/ComplexDialect.cpp  | 10 +---------
 2 files changed, 18 insertions(+), 9 deletions(-)

diff --git a/mlir/include/mlir/Dialect/Complex/IR/ComplexAttributes.td b/mlir/include/mlir/Dialect/Complex/IR/ComplexAttributes.td
index c5e099a83719c..ab6074eccb538 100644
--- a/mlir/include/mlir/Dialect/Complex/IR/ComplexAttributes.td
+++ b/mlir/include/mlir/Dialect/Complex/IR/ComplexAttributes.td
@@ -37,8 +37,25 @@ def Complex_NumberAttr : Complex_Attr<"Number", "number"> {
   let parameters = (ins APFloatParameter<"">:$real,
                         APFloatParameter<"">:$imag,
                         AttributeSelfTypeParameter<"">:$type);
+  let builders = [
+    AttrBuilderWithInferredContext<(ins "mlir::ComplexType":$type,
+                                        "double":$real,
+                                        "double":$imag), [{
+      auto elementType = type.getElementType().cast<FloatType>();
+      APFloat realFloat(real);
+      bool unused;
+      realFloat.convert(elementType.getFloatSemantics(),
+                        APFloat::rmNearestTiesToEven, &unused);
+      APFloat imagFloat(imag);
+      imagFloat.convert(elementType.getFloatSemantics(),
+                        APFloat::rmNearestTiesToEven, &unused);
+      return $_get(type.getContext(), realFloat, imagFloat, type);
+    }]>
+  ];
+
   let genVerifyDecl = 1;
   let hasCustomAssemblyFormat = 1;
+  let skipDefaultBuilders = 1;
 }
 
 #endif // COMPLEX_ATTRIBUTE
diff --git a/mlir/lib/Dialect/Complex/IR/ComplexDialect.cpp b/mlir/lib/Dialect/Complex/IR/ComplexDialect.cpp
index 8c5308929659d..c93d8811712f8 100644
--- a/mlir/lib/Dialect/Complex/IR/ComplexDialect.cpp
+++ b/mlir/lib/Dialect/Complex/IR/ComplexDialect.cpp
@@ -81,13 +81,5 @@ Attribute complex::NumberAttr::parse(AsmParser &parser, Type odsType) {
       parser.parseFloat(imag) || parser.parseGreater())
     return {};
 
-  bool unused = false;
-  APFloat realFloat(real);
-  realFloat.convert(type.cast<FloatType>().getFloatSemantics(),
-                    APFloat::rmNearestTiesToEven, &unused);
-  APFloat imagFloat(imag);
-  imagFloat.convert(type.cast<FloatType>().getFloatSemantics(),
-                    APFloat::rmNearestTiesToEven, &unused);
-  return NumberAttr::get(parser.getContext(), realFloat, imagFloat,
-                         ComplexType::get(type));
+  return NumberAttr::get(ComplexType::get(type), real, imag);
 }