[mlir][sparse] fix bug in complex zero detection

[aartbik]

Summary

We were checking real-part twice, not real/imag-part.
The new test only passes after the bug fix.

mlir/lib/Dialect/SparseTensor/Utils/Merger.cpp

 }
 
 /// Only returns false if we are certain this is a nonzero.
 bool Merger::maybeZero(unsigned e) const {
   if (tensorExps[e].kind == kInvariant) {
     if (auto c = tensorExps[e].val.getDefiningOp<complex::ConstantOp>()) {
       ArrayAttr arrayAttr = c.getValue();
       return arrayAttr[0].cast<FloatAttr>().getValue().isZero() &&
-             arrayAttr[0].cast<FloatAttr>().getValue().isZero();
+             arrayAttr[1].cast<FloatAttr>().getValue().isZero();
     }
     if (auto c = tensorExps[e].val.getDefiningOp<arith::ConstantIntOp>())
       return c.value() == 0;
     if (auto c = tensorExps[e].val.getDefiningOp<arith::ConstantFloatOp>())
       return c.value().isZero();
   }
   return true;
 }

mlir/test/Dialect/SparseTensor/sparse_fp_ops.mlir

-// NOTE: Assertions have been autogenerated by utils/generate-test-checks.py
 // RUN: mlir-opt %s -sparsification | FileCheck %s
 
 #SV = #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>
 
 #trait1 = {
   indexing_maps = [
     affine_map<(i) -> (i)>,  // a
     affine_map<(i) -> (i)>   // x (out)
%0 = math.abs %a : f64
         %4 = math.expm1 %3 : f64
         %5 = math.log1p %4 : f64
         %6 = math.sin %5 : f64
         %7 = math.tanh %6 : f64
         linalg.yield %7 : f64
   } -> tensor<32xf64, #SV>
   return %0 : tensor<32xf64, #SV>
 }
+
+// CHECK-LABEL: func.func @complex_divbyc(
+// CHECK-SAME:    %[[VAL_0:.*]]: tensor<32xcomplex<f64>, #sparse_tensor.encoding<{{.*}}>> {
+// CHECK-DAG:     %[[VAL_1:.*]] = arith.constant 0 : index
+// CHECK-DAG:     %[[VAL_2:.*]] = arith.constant 1 : index
+// CHECK-DAG:     %[[VAL_3:.*]] = complex.constant [0.000000e+00, 1.000000e+00] : complex<f64>
+// CHECK:         %[[VAL_4:.*]] = bufferization.alloc_tensor() : tensor<32xcomplex<f64>, #sparse_tensor.encoding<{{.*}}>>
+// CHECK:         %[[VAL_5:.*]] = sparse_tensor.pointers %[[VAL_0]], %[[VAL_1]] : tensor<32xcomplex<f64>, #sparse_tensor.encoding<{{.*}}>> to memref<?xindex>
+// CHECK:         %[[VAL_6:.*]] = sparse_tensor.indices %[[VAL_0]], %[[VAL_1]] : tensor<32xcomplex<f64>, #sparse_tensor.encoding<{{.*}}>> to memref<?xindex>
+// CHECK:         %[[VAL_7:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32xcomplex<f64>, #sparse_tensor.encoding<{{.*}}>> to memref<?xcomplex<f64>>
+// CHECK:         %[[VAL_8:.*]] = memref.alloca(%[[VAL_2]]) : memref<?xindex>
+// CHECK:         %[[VAL_9:.*]] = memref.alloca() : memref<complex<f64>>
+// CHECK:         %[[VAL_10:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_1]]] : memref<?xindex>
+// CHECK:         %[[VAL_11:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_2]]] : memref<?xindex>
+// CHECK:         scf.for %[[VAL_12:.*]] = %[[VAL_10]] to %[[VAL_11]] step %[[VAL_2]] {
+// CHECK:           %[[VAL_13:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_12]]] : memref<?xindex>
+// CHECK:           memref.store %[[VAL_13]], %[[VAL_8]]{{\[}}%[[VAL_1]]] : memref<?xindex>
+// CHECK:           %[[VAL_14:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_12]]] : memref<?xcomplex<f64>>
+// CHECK:           %[[VAL_15:.*]] = complex.div %[[VAL_14]], %[[VAL_3]] : complex<f64>
+// CHECK:           memref.store %[[VAL_15]], %[[VAL_9]][] : memref<complex<f64>>
+// CHECK:           sparse_tensor.lex_insert %[[VAL_4]], %[[VAL_8]], %[[VAL_9]] : tensor<32xcomplex<f64>, #sparse_tensor.encoding<{{.*}}>>, memref<?xindex>, memref<complex<f64>>
+// CHECK:         }
+// CHECK:         %[[VAL_16:.*]] = sparse_tensor.load %[[VAL_4]] hasInserts : tensor<32xcomplex<f64>, #sparse_tensor.encoding<{{.*}}>>
+// CHECK:         return %[[VAL_16]] : tensor<32xcomplex<f64>, #sparse_tensor.encoding<{{.*}}>>
+// CHECK:       }
+func.func @complex_divbyc(%arg0: tensor<32xcomplex<f64>, #SV>) -> tensor<32xcomplex<f64>, #SV> {
+  %c = complex.constant [0.0, 1.0] : complex<f64>
+  %init = bufferization.alloc_tensor() : tensor<32xcomplex<f64>, #SV>
+  %0 = linalg.generic #traitc
+     ins(%arg0: tensor<32xcomplex<f64>, #SV>)
+    outs(%init: tensor<32xcomplex<f64>, #SV>) {
+      ^bb(%a: complex<f64>, %x: complex<f64>):
+        %0 = complex.div %a, %c : complex<f64>
+        linalg.yield %0 : complex<f64>
+  } -> tensor<32xcomplex<f64>, #SV>
+  return %0 : tensor<32xcomplex<f64>, #SV>
+}