[mlir][sparse] Add more complex operations.

[bixia]

Summary

Support complex operations sqrt, expm1, and tanh.

Add tests.

[aartbik]

Onward to feature completeness!

mlir/include/mlir/Dialect/SparseTensor/Utils/Merger.h

enum Kind {
   kInvariant,
   kIndex,
   // Unary operations.
   kAbsF,
   kAbsC,
   kCeilF,
   kFloorF,
   kSqrtF,
+  kSqrtC,
   kExpm1F,
+  kExpm1C,
   kLog1pF,
   kLog1pC,
   kSinF,
   kSinC,
   kTanhF,
+  kTanhC,
   kNegF,
   kNegC,
   kNegI,
   kTruncF,
   kExtF,
   kCastFS, // signed
   kCastFU, // unsigned
   kCastSF, // signed

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

case kIndex:
     assert(x != -1u && y == -1u && !v && !o);
     index = x;
     break;
   case kAbsF:
   case kAbsC:
   case kCeilF:
   case kFloorF:
   case kSqrtF:
+  case kSqrtC:
   case kExpm1F:
+  case kExpm1C:
   case kLog1pF:
   case kLog1pC:
   case kSinF:
   case kSinC:
   case kTanhF:
+  case kTanhC:
   case kNegF:
   case kNegC:
   case kNegI:
   case kCIm:
   case kCRe:
     assert(x != -1u && y == -1u && !v && !o);
     children.e0 = x;
     children.e1 = y;
bool Merger::isSingleCondition(unsigned t, unsigned e) const {
   switch (tensorExps[e].kind) {
   case kTensor:
     return tensorExps[e].tensor == t;
   case kAbsF:
   case kAbsC:
   case kCeilF:
   case kFloorF:
   case kSqrtF:
+  case kSqrtC:
   case kExpm1F:
+  case kExpm1C:
   case kLog1pF:
   case kLog1pC:
   case kSinF:
   case kSinC:
   case kTanhF:
+  case kTanhC:
   case kNegF:
   case kNegC:
   case kNegI:
   case kTruncF:
   case kExtF:
   case kCastFS:
   case kCastFU:
   case kCastSF:
static const char *kindToOpSymbol(Kind kind) {
   case kAbsF:
   case kAbsC:
     return "abs";
   case kCeilF:
     return "ceil";
   case kFloorF:
     return "floor";
   case kSqrtF:
+  case kSqrtC:
     return "sqrt";
   case kExpm1F:
+  case kExpm1C:
     return "expm1";
   case kLog1pF:
   case kLog1pC:
     return "log1p";
   case kSinF:
   case kSinC:
     return "sin";
   case kTanhF:
+  case kTanhC:
     return "tanh";
   case kNegF:
   case kNegC:
   case kNegI:
     return "-";
   case kTruncF:
   case kExtF:
   case kCastFS:
case kInvariant:
     break;
   case kIndex:
     llvm::dbgs() << "index_" << tensorExps[e].index;
     break;
   case kAbsF:
   case kCeilF:
   case kFloorF:
   case kSqrtF:
+  case kSqrtC:
   case kExpm1F:
+  case kExpm1C:
   case kLog1pF:
   case kSinF:
   case kTanhF:
+  case kTanhC:
   case kNegF:
   case kNegI:
   case kTruncF:
   case kExtF:
   case kCastFS:
   case kCastFU:
   case kCastSF:
   case kCastUF:
unsigned Merger::buildLattices(unsigned e, unsigned i) {
   }
   case kAbsF:
   case kAbsC:
   case kCeilF:
   case kCIm:
   case kCRe:
   case kFloorF:
   case kSqrtF:
+  case kSqrtC:
   case kExpm1F:
+  case kExpm1C:
   case kLog1pF:
   case kLog1pC:
   case kSinF:
   case kSinC:
   case kTanhF:
+  case kTanhC:
   case kNegF:
   case kNegC:
   case kNegI:
   case kTruncF:
   case kExtF:
   case kCastFS:
   case kCastFU:
   case kCastSF:
if (x.hasValue()) {
       if (isa<complex::AbsOp>(def))
         return addExp(kAbsC, e);
       if (isa<math::CeilOp>(def))
         return addExp(kCeilF, e);
       if (isa<math::FloorOp>(def))
         return addExp(kFloorF, e);
       if (isa<math::SqrtOp>(def))
         return addExp(kSqrtF, e);
+      if (isa<complex::SqrtOp>(def))
+        return addExp(kSqrtC, e);
       if (isa<math::ExpM1Op>(def))
         return addExp(kExpm1F, e);
+      if (isa<complex::Expm1Op>(def))
+        return addExp(kExpm1C, e);
       if (isa<math::Log1pOp>(def))
         return addExp(kLog1pF, e);
       if (isa<complex::Log1pOp>(def))
         return addExp(kLog1pC, e);
       if (isa<math::SinOp>(def))
         return addExp(kSinF, e);
       if (isa<complex::SinOp>(def))
         return addExp(kSinC, e);
       if (isa<math::TanhOp>(def))
         return addExp(kTanhF, e);
+      if (isa<complex::TanhOp>(def))
+        return addExp(kTanhC, e);
       if (isa<arith::NegFOp>(def))
         return addExp(kNegF, e); // no negi in std
       if (isa<complex::NegOp>(def))
         return addExp(kNegC, e);
       if (isa<arith::TruncFOp>(def))
         return addExp(kTruncF, e, v);
       if (isa<arith::ExtFOp>(def))
         return addExp(kExtF, e, v);
case kAbsC: {
     return rewriter.create<complex::AbsOp>(loc, eltType, v0);
   }
   case kCeilF:
     return rewriter.create<math::CeilOp>(loc, v0);
   case kFloorF:
     return rewriter.create<math::FloorOp>(loc, v0);
   case kSqrtF:
     return rewriter.create<math::SqrtOp>(loc, v0);
+  case kSqrtC:
+    return rewriter.create<complex::SqrtOp>(loc, v0);
   case kExpm1F:
     return rewriter.create<math::ExpM1Op>(loc, v0);
+  case kExpm1C:
+    return rewriter.create<complex::Expm1Op>(loc, v0);
   case kLog1pF:
     return rewriter.create<math::Log1pOp>(loc, v0);
   case kLog1pC:
     return rewriter.create<complex::Log1pOp>(loc, v0);
   case kSinF:
     return rewriter.create<math::SinOp>(loc, v0);
   case kSinC:
     return rewriter.create<complex::SinOp>(loc, v0);
   case kTanhF:
     return rewriter.create<math::TanhOp>(loc, v0);
+  case kTanhC:
+    return rewriter.create<complex::TanhOp>(loc, v0);
   case kNegF:
     return rewriter.create<arith::NegFOp>(loc, v0);
   case kNegC:
     return rewriter.create<complex::NegOp>(loc, v0);
   case kNegI: // no negi in std
     return rewriter.create<arith::SubIOp>(
         loc,
         rewriter.create<arith::ConstantOp>(loc, v0.getType(),

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_complex_ops.mlir

%0 = linalg.generic #trait_op1
         outs(%xv: tensor<?xcomplex<f64>, #SparseVector>) {
         ^bb(%a: complex<f64>, %x: complex<f64>):
           %1 = complex.sin %a : complex<f64>
           linalg.yield %1 : complex<f64>
     } -> tensor<?xcomplex<f64>, #SparseVector>
     return %0 : tensor<?xcomplex<f64>, #SparseVector>
   }
 
+  func.func @complex_sqrt(%arga: tensor<?xcomplex<f64>, #SparseVector>)
+                 -> tensor<?xcomplex<f64>, #SparseVector> {
+    %c0 = arith.constant 0 : index
+    %d = tensor.dim %arga, %c0 : tensor<?xcomplex<f64>, #SparseVector>
+    %xv = sparse_tensor.init [%d] : tensor<?xcomplex<f64>, #SparseVector>
+    %0 = linalg.generic #trait_op1
+       ins(%arga: tensor<?xcomplex<f64>, #SparseVector>)
+        outs(%xv: tensor<?xcomplex<f64>, #SparseVector>) {
+        ^bb(%a: complex<f64>, %x: complex<f64>):
+          %1 = complex.sqrt %a : complex<f64>
+          linalg.yield %1 : complex<f64>
+    } -> tensor<?xcomplex<f64>, #SparseVector>
+    return %0 : tensor<?xcomplex<f64>, #SparseVector>
+  }
+
+  func.func @complex_tanh(%arga: tensor<?xcomplex<f64>, #SparseVector>)
+                 -> tensor<?xcomplex<f64>, #SparseVector> {
+    %c0 = arith.constant 0 : index
+    %d = tensor.dim %arga, %c0 : tensor<?xcomplex<f64>, #SparseVector>
+    %xv = sparse_tensor.init [%d] : tensor<?xcomplex<f64>, #SparseVector>
+    %0 = linalg.generic #trait_op1
+       ins(%arga: tensor<?xcomplex<f64>, #SparseVector>)
+        outs(%xv: tensor<?xcomplex<f64>, #SparseVector>) {
+       ^bb(%a: complex<f64>, %x: complex<f64>):
+          %1 = complex.tanh %a : complex<f64>
+          linalg.yield %1 : complex<f64>
+   } -> tensor<?xcomplex<f64>, #SparseVector>
+    return %0 : tensor<?xcomplex<f64>, #SparseVector>
+  }
+
+  func.func @clog1p_expm1(%arga: tensor<?xcomplex<f64>, #SparseVector>)
+                 -> tensor<?xcomplex<f64>, #SparseVector> {
+    %c0 = arith.constant 0 : index
+    %d = tensor.dim %arga, %c0 : tensor<?xcomplex<f64>, #SparseVector>
+    %xv = sparse_tensor.init [%d] : tensor<?xcomplex<f64>, #SparseVector>
+    %0 = linalg.generic #trait_op1
+       ins(%arga: tensor<?xcomplex<f64>, #SparseVector>)
+        outs(%xv: tensor<?xcomplex<f64>, #SparseVector>) {
+        ^bb(%a: complex<f64>, %x: complex<f64>):
+          %1 = complex.log1p %a : complex<f64>
+          // TODO(bixia): Enable this line after adding complex.expm1 to
+          // complex to standard lowering.
+          // %2 = complex.expm1 %1 : complex<f64>
+          linalg.yield %1 : complex<f64>
+    } -> tensor<?xcomplex<f64>, #SparseVector>
+    return %0 : tensor<?xcomplex<f64>, #SparseVector>
+  }
+
   func.func @cdiv(%arga: tensor<?xcomplex<f64>, #SparseVector>)
                  -> tensor<?xcomplex<f64>, #SparseVector> {
     %c0 = arith.constant 0 : index
     %d = tensor.dim %arga, %c0 : tensor<?xcomplex<f64>, #SparseVector>
     %xv = sparse_tensor.init [%d] : tensor<?xcomplex<f64>, #SparseVector>
     %c = complex.constant [2.0 : f64, 0.0 : f64] : complex<f64>
     %0 = linalg.generic #trait_op1
        ins(%arga: tensor<?xcomplex<f64>, #SparseVector>)
func.func @entry() {
     %sv2 = sparse_tensor.convert %v2 : tensor<32xcomplex<f64>> to tensor<?xcomplex<f64>, #SparseVector>
 
     // Call sparse vector kernels.
     %0 = call @cops(%sv1, %sv2)
        : (tensor<?xcomplex<f64>, #SparseVector>,
           tensor<?xcomplex<f64>, #SparseVector>) -> tensor<?xcomplex<f64>, #SparseVector>
     %1 = call @csin(%sv1)
        : (tensor<?xcomplex<f64>, #SparseVector>) -> tensor<?xcomplex<f64>, #SparseVector>
-    %2 = call @cdiv(%sv1)
+    %2 = call @complex_sqrt(%sv1)
+       : (tensor<?xcomplex<f64>, #SparseVector>) -> tensor<?xcomplex<f64>, #SparseVector>
+    %3 = call @complex_tanh(%sv2)
        : (tensor<?xcomplex<f64>, #SparseVector>) -> tensor<?xcomplex<f64>, #SparseVector>
-    %3 = call @cabs(%sv1)
+    %4 = call @clog1p_expm1(%sv1)
+       : (tensor<?xcomplex<f64>, #SparseVector>) -> tensor<?xcomplex<f64>, #SparseVector>
+    %5 = call @cdiv(%sv1)
+       : (tensor<?xcomplex<f64>, #SparseVector>) -> tensor<?xcomplex<f64>, #SparseVector>
+    %6 = call @cabs(%sv1)
        : (tensor<?xcomplex<f64>, #SparseVector>) -> tensor<?xf64, #SparseVector>
 
     //
     // Verify the results.
     //
     %d3 = arith.constant 3 : index
     %d4 = arith.constant 4 : index
     // CHECK: -5.13
     // CHECK-NEXT: 2
     // CHECK-NEXT: 1
     // CHECK-NEXT: 0
     // CHECK-NEXT: 1
     // CHECK-NEXT: 4
     // CHECK-NEXT: 8
     // CHECK-NEXT: 6
     call @dumpc(%0, %d4) : (tensor<?xcomplex<f64>, #SparseVector>, index) -> ()
     // CHECK-NEXT: 3.43887
     // CHECK-NEXT: 1.47097
     // CHECK-NEXT: 3.85374
     // CHECK-NEXT: -27.0168
     // CHECK-NEXT: -193.43
     // CHECK-NEXT: 57.2184
     call @dumpc(%1, %d3) : (tensor<?xcomplex<f64>, #SparseVector>, index) -> ()
+    // CHECK-NEXT: 0.433635
+    // CHECK-NEXT: 2.30609
+    // CHECK-NEXT: 2
+    // CHECK-NEXT: 1
+    // CHECK-NEXT: 2.53083
+    // CHECK-NEXT: 1.18538
+    call @dumpc(%2, %d3) : (tensor<?xcomplex<f64>, #SparseVector>, index) -> ()
+    // CHECK-NEXT: 0.761594
+    // CHECK-NEXT: 0
+    // CHECK-NEXT: -0.964028
+    // CHECK-NEXT: 0
+    // CHECK-NEXT: 0.995055
+    // CHECK-NEXT: 0
+    call @dumpc(%3, %d3) : (tensor<?xcomplex<f64>, #SparseVector>, index) -> ()
+    // CHECK-NEXT: 1.52361
+    // CHECK-NEXT: 2.69061
+    // CHECK-NEXT: 1.73287
+    // CHECK-NEXT: 0.785398
+    // CHECK-NEXT: 2.13833
+    // CHECK-NEXT: 0.785398
+    call @dumpc(%4, %d3) : (tensor<?xcomplex<f64>, #SparseVector>, index) -> ()
     // CHECK-NEXT: -2.565
     // CHECK-NEXT: 1
     // CHECK-NEXT: 1.5
     // CHECK-NEXT: 2
     // CHECK-NEXT: 2.5
     // CHECK-NEXT: 3
-    call @dumpc(%2, %d3) : (tensor<?xcomplex<f64>, #SparseVector>, index) -> ()
+    call @dumpc(%5, %d3) : (tensor<?xcomplex<f64>, #SparseVector>, index) -> ()
     // CHECK-NEXT: ( 5.50608, 5, 7.81025 )
-    call @dumpf(%3) : (tensor<?xf64, #SparseVector>) -> ()
+    call @dumpf(%6) : (tensor<?xf64, #SparseVector>) -> ()
 
     // Release the resources.
     sparse_tensor.release %sv1 : tensor<?xcomplex<f64>, #SparseVector>
     sparse_tensor.release %sv2 : tensor<?xcomplex<f64>, #SparseVector>
     sparse_tensor.release %0 : tensor<?xcomplex<f64>, #SparseVector>
     sparse_tensor.release %1 : tensor<?xcomplex<f64>, #SparseVector>
     sparse_tensor.release %2 : tensor<?xcomplex<f64>, #SparseVector>
-    sparse_tensor.release %3 : tensor<?xf64, #SparseVector>
+    sparse_tensor.release %3 : tensor<?xcomplex<f64>, #SparseVector>
+    sparse_tensor.release %4 : tensor<?xcomplex<f64>, #SparseVector>
+    sparse_tensor.release %5 : tensor<?xcomplex<f64>, #SparseVector>
+    sparse_tensor.release %6 : tensor<?xf64, #SparseVector>
     return
   }
 }