[Mlir-commits] [mlir] [MLIR][Complex] DivOp check for NaN in the folder (PR #169724)

[Amr Hesham]

================
@@ -374,21 +374,30 @@ OpFoldResult MulOp::fold(FoldAdaptor adaptor) {
 
 OpFoldResult DivOp::fold(FoldAdaptor adaptor) {
   auto rhs = adaptor.getRhs();
-  if (!rhs)
+  auto lhs = adaptor.getLhs();
+  if (!rhs || !lhs)
     return {};
 
-  ArrayAttr arrayAttr = dyn_cast<ArrayAttr>(rhs);
-  if (!arrayAttr || arrayAttr.size() != 2)
+  ArrayAttr rhsArrayAttr = dyn_cast<ArrayAttr>(rhs);
+  if (!rhsArrayAttr || rhsArrayAttr.size() != 2)
     return {};
 
-  APFloat real = cast<FloatAttr>(arrayAttr[0]).getValue();
-  APFloat imag = cast<FloatAttr>(arrayAttr[1]).getValue();
+  ArrayAttr lhsArrayAttr = dyn_cast<ArrayAttr>(lhs);
+  if (!lhsArrayAttr || lhsArrayAttr.size() != 2)
+    return {};
 
-  if (!imag.isZero())
+  APFloat rhsImag = cast<FloatAttr>(rhsArrayAttr[1]).getValue();
+  if (!rhsImag.isZero())
+    return {};
+
+  APFloat lhsReal = cast<FloatAttr>(lhsArrayAttr[0]).getValue();
+  APFloat lhsImag = cast<FloatAttr>(lhsArrayAttr[1]).getValue();
+  if (lhsReal.isNaN() || lhsImag.isNaN())
     return {};
----------------
AmrDeveloper wrote:

@jpienaar Thanks, it's a bit cleaner now

https://github.com/llvm/llvm-project/pull/169724