[mlir][sparse] Add support for complex.im and complex.re to the sparse compiler.

[bixia]

Summary

Add a test.

[aartbik]

Nice! Very fast. A few nits, but good to go.

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

35

This is very nitpicky, but I would, as far as enum order goes, put these with the casts below, for example, right before kBitCast (since I see these as converting complex to float). Right now, they break the abs/ceil/floor/sqrt sequence a bit. If you agree, then make sure they appear at the same place in switches.

(note that it is more grouped by function, and less so alphabetically)

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

910

eltType?

also, we need to verify how (and if!) this interacts with vectorization

[bixia]

Relocated the new enum values.

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

35

Moved the two enums to right before kBitCast, also reorder other places where they are used.

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

 enum Kind {
   // Leaf.
   kTensor = 0,
   kInvariant,
   kIndex,
   // Unary operations.
   kAbsF,
   kCeilF,
   kFloorF,
   kSqrtF,
   kExpm1F,
   kLog1pF,
   kSinF,
   kTanhF,
   kNegF,
   kNegI,
   kTruncF,
   kExtF,
   kCastFS, // signed
   kCastFU, // unsigned
   kCastSF, // signed
   kCastUF, // unsigned
   kCastS,  // signed
   kCastU,  // unsigned
   kCastIdx,
   kTruncI,
+  kCIm, // complex.im
+  kCRe, // complex.re
   kBitCast,
   kBinaryBranch, // semiring unary branch created from a binary op
   kUnary,        // semiring unary op
   // Binary operations.
   kMulF,
   kMulC,
   kMulI,
   kDivF,

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

TensorExp::TensorExp(Kind k, unsigned x, unsigned y, Value v, Operation *o)
   case kFloorF:
   case kSqrtF:
   case kExpm1F:
   case kLog1pF:
   case kSinF:
   case kTanhF:
   case kNegF:
   case kNegI:
+  case kCIm:
+  case kCRe:
     assert(x != -1u && y == -1u && !v && !o);
     children.e0 = x;
     children.e1 = y;
     break;
   case kTruncF:
   case kExtF:
   case kCastFS:
   case kCastFU:
bool Merger::isSingleCondition(unsigned t, unsigned e) const {
   case kCastFS:
   case kCastFU:
   case kCastSF:
   case kCastUF:
   case kCastS:
   case kCastU:
   case kCastIdx:
   case kTruncI:
+  case kCIm:
+  case kCRe:
   case kBitCast:
     return isSingleCondition(t, tensorExps[e].children.e0);
   case kDivF: // note: x / c only
   case kDivS:
   case kDivU:
     assert(!maybeZero(tensorExps[e].children.e1));
     return isSingleCondition(t, tensorExps[e].children.e0);
   case kShrS: // note: x >> inv only
static const char *kindToOpSymbol(Kind kind) {
   case kCastFS:
   case kCastFU:
   case kCastSF:
   case kCastUF:
   case kCastS:
   case kCastU:
   case kCastIdx:
   case kTruncI:
+  case kCIm:
+    return "complex.im";
+  case kCRe:
+    return "complex.re";
   case kBitCast:
     return "cast";
   case kBinaryBranch:
     return "binary_branch";
   case kUnary:
     return "unary";
   case kMulF:
   case kMulC:
if (kind == kTensor) {
       if (hasSparseOut && t == outTensor)
         t = syntheticTensor;
     }
     latSets[s].push_back(addLat(t, i, e));
     return s;
   }
   case kAbsF:
   case kCeilF:
+  case kCIm:
+  case kCRe:
   case kFloorF:
   case kSqrtF:
   case kExpm1F:
   case kLog1pF:
   case kSinF:
   case kTanhF:
   case kNegF:
   case kNegI:
if (x.hasValue()) {
       if (isa<arith::ExtSIOp>(def))
         return addExp(kCastS, e, v);
       if (isa<arith::ExtUIOp>(def))
         return addExp(kCastU, e, v);
       if (isa<arith::IndexCastOp>(def))
         return addExp(kCastIdx, e, v);
       if (isa<arith::TruncIOp>(def))
         return addExp(kTruncI, e, v);
+      if (isa<complex::ImOp>(def))
+        return addExp(kCIm, e);
+      if (isa<complex::ReOp>(def))
+        return addExp(kCRe, e);
       if (isa<arith::BitcastOp>(def))
         return addExp(kBitCast, e, v);
       if (isa<sparse_tensor::UnaryOp>(def))
         return addExp(kUnary, e, Value(), def);
     }
   }
   // Construct binary operations if subexpressions can be built.
   // See buildLattices() for an explanation of rejecting certain
Value Merger::buildExp(RewriterBase &rewriter, Location loc, unsigned e,
   // Unary ops.
   case kAbsF:
     return rewriter.create<math::AbsOp>(loc, 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 kExpm1F:
     return rewriter.create<math::ExpM1Op>(loc, v0);
   case kLog1pF:
     return rewriter.create<math::Log1pOp>(loc, v0);
   case kSinF:
     return rewriter.create<math::SinOp>(loc, v0);
   case kTanhF:
     return rewriter.create<math::TanhOp>(loc, v0);
Value Merger::buildExp(RewriterBase &rewriter, Location loc, unsigned e,
   case kCastS:
     return rewriter.create<arith::ExtSIOp>(loc, inferType(e, v0), v0);
   case kCastU:
     return rewriter.create<arith::ExtUIOp>(loc, inferType(e, v0), v0);
   case kCastIdx:
     return rewriter.create<arith::IndexCastOp>(loc, inferType(e, v0), v0);
   case kTruncI:
     return rewriter.create<arith::TruncIOp>(loc, inferType(e, v0), v0);
+  case kCIm:
+  case kCRe: {
+    auto type = v0.getType().template cast<ComplexType>();
+    auto eltType = type.getElementType().template cast<FloatType>();
+    if (tensorExps[e].kind == kCIm)
+      return rewriter.create<complex::ImOp>(loc, eltType, v0);
+
+    return rewriter.create<complex::ReOp>(loc, eltType, v0);
+  }
   case kBitCast:
     return rewriter.create<arith::BitcastOp>(loc, inferType(e, v0), v0);
   // Binary ops.
   case kMulF:
     return rewriter.create<arith::MulFOp>(loc, v0, v1);
   case kMulC:
     return rewriter.create<complex::MulOp>(loc, v0, v1);
   case kMulI:

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

+// RUN: mlir-opt %s --sparse-compiler | \
+// RUN: mlir-cpu-runner \
+// RUN:  -e entry -entry-point-result=void  \
+// RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
+// RUN: FileCheck %s
+
+#SparseVector = #sparse_tensor.encoding<{dimLevelType = ["compressed"]}>
+
+#trait_op = {
+  indexing_maps = [
+    affine_map<(i) -> (i)>,  // a (in)
+    affine_map<(i) -> (i)>   // x (out)
+  ],
+  iterator_types = ["parallel"],
+  doc = "x(i) = OP a(i)"
+}
+
+module {
+  func.func @cre(%arga: tensor<?xcomplex<f32>, #SparseVector>)
+                -> tensor<?xf32, #SparseVector> {
+    %c = arith.constant 0 : index
+    %d = tensor.dim %arga, %c : tensor<?xcomplex<f32>, #SparseVector>
+    %xv = sparse_tensor.init [%d] : tensor<?xf32, #SparseVector>
+    %0 = linalg.generic #trait_op
+       ins(%arga: tensor<?xcomplex<f32>, #SparseVector>)
+        outs(%xv: tensor<?xf32, #SparseVector>) {
+        ^bb(%a: complex<f32>, %x: f32):
+          %1 = complex.re %a : complex<f32>
+          linalg.yield %1 : f32
+    } -> tensor<?xf32, #SparseVector>
+    return %0 : tensor<?xf32, #SparseVector>
+  }
+
+  func.func @cim(%arga: tensor<?xcomplex<f32>, #SparseVector>)
+                -> tensor<?xf32, #SparseVector> {
+    %c = arith.constant 0 : index
+    %d = tensor.dim %arga, %c : tensor<?xcomplex<f32>, #SparseVector>
+    %xv = sparse_tensor.init [%d] : tensor<?xf32, #SparseVector>
+    %0 = linalg.generic #trait_op
+       ins(%arga: tensor<?xcomplex<f32>, #SparseVector>)
+        outs(%xv: tensor<?xf32, #SparseVector>) {
+        ^bb(%a: complex<f32>, %x: f32):
+          %1 = complex.im %a : complex<f32>
+          linalg.yield %1 : f32
+    } -> tensor<?xf32, #SparseVector>
+    return %0 : tensor<?xf32, #SparseVector>
+  }
+
+  func.func @dump(%arg0: tensor<?xf32, #SparseVector>) {
+    %c0 = arith.constant 0 : index
+    %d0 = arith.constant -1.0 : f32
+    %values = sparse_tensor.values %arg0 : tensor<?xf32, #SparseVector> to memref<?xf32>
+    %0 = vector.transfer_read %values[%c0], %d0: memref<?xf32>, vector<4xf32>
+    vector.print %0 : vector<4xf32>
+    %indices = sparse_tensor.indices %arg0, %c0 : tensor<?xf32, #SparseVector> to memref<?xindex>
+    %1 = vector.transfer_read %indices[%c0], %c0: memref<?xindex>, vector<4xindex>
+    vector.print %1 : vector<4xindex>
+    return
+  }
+
+  // Driver method to call and verify functions cim and cre.
+  func.func @entry() {
+    // Setup sparse vectors.
+    %v1 = arith.constant sparse<
+       [ [0], [20], [31] ],
+         [ (5.13, 2.0), (3.0, 4.0), (5.0, 6.0) ] > : tensor<32xcomplex<f32>>
+    %sv1 = sparse_tensor.convert %v1 : tensor<32xcomplex<f32>> to tensor<?xcomplex<f32>, #SparseVector>
+
+    // Call sparse vector kernels.
+    %0 = call @cre(%sv1)
+       : (tensor<?xcomplex<f32>, #SparseVector>) -> tensor<?xf32, #SparseVector>
+
+    %1 = call @cim(%sv1)
+       : (tensor<?xcomplex<f32>, #SparseVector>) -> tensor<?xf32, #SparseVector>
+
+    //
+    // Verify the results.
+    //
+    // CHECK: ( 5.13, 3, 5, -1 )
+    // CHECK-NEXT: ( 0, 20, 31, 0 )
+    // CHECK-NEXT: ( 2, 4, 6, -1 )
+    // CHECK-NEXT: ( 0, 20, 31, 0 )
+    //
+    call @dump(%0) : (tensor<?xf32, #SparseVector>) -> ()
+    call @dump(%1) : (tensor<?xf32, #SparseVector>) -> ()
+
+    // Release the resources.
+    sparse_tensor.release %sv1 : tensor<?xcomplex<f32>, #SparseVector>
+    sparse_tensor.release %0 : tensor<?xf32, #SparseVector>
+    sparse_tensor.release %1 : tensor<?xf32, #SparseVector>
+    return
+  }
+}