[mlir][sparse] introduce complex type to sparse tensor support

[aartbik]

Summary

This is the first implementation of complex (f64 and f32) support
in the sparse compiler, with complex add/mul as first operations.
Note that various features are still TBD, such as other ops, and
reading in complex values from file. Also, note that the
std::complex<float> had a bit of an ABI issue when passed as
single argument. It is still TBD if better solutions are possible.

[aartbik]

CMake fix

[bixia]

mlir/lib/ExecutionEngine/SparseTensorUtils.cpp

1355

It would be nice if we describe the problem we are trying to workaround here.

1381

Similar to the TODO above.

[aartbik]

mlir/lib/ExecutionEngine/SparseTensorUtils.cpp

1355

Good suggestion, added a comment.

[aartbik]

addressed comments

[aartbik]

correct include for complex type

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

enum Kind {
   kCastU,  // unsigned
   kCastIdx,
   kTruncI,
   kBitCast,
   kBinaryBranch, // semiring unary branch created from a binary op
   kUnary,        // semiring unary op
   // Binary operations.
   kMulF,
+  kMulC,
   kMulI,
   kDivF,
   kDivS, // signed
   kDivU, // unsigned
   kAddF,
+  kAddC,
   kAddI,
   kSubF,
   kSubI,
   kAndI,
   kOrI,
   kXorI,
   kShrS, // signed
   kShrU, // unsigned

mlir/include/mlir/ExecutionEngine/SparseTensorUtils.h

 
 /// Encoding of the elemental type, for "overloading" @newSparseTensor.
 enum class PrimaryType : uint32_t {
   kF64 = 1,
   kF32 = 2,
   kI64 = 3,
   kI32 = 4,
   kI16 = 5,
-  kI8 = 6
+  kI8 = 6,
+  kC64 = 7,
+  kC32 = 8
 };
 
 /// The actions performed by @newSparseTensor.
 enum class Action : uint32_t {
   kEmpty = 0,
   kFromFile = 1,
   kFromCOO = 2,
   kSparseToSparse = 3,

mlir/lib/Dialect/SparseTensor/Pipelines/CMakeLists.txt

 add_mlir_dialect_library(MLIRSparseTensorPipelines
   SparseTensorPipelines.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/SparseTensor
 
   LINK_LIBS PUBLIC
   MLIRArithmeticTransforms
   MLIRAffineToStandard
   MLIRBufferizationTransforms
+  MLIRComplexToLLVM
+  MLIRComplexToStandard
   MLIRFuncTransforms
   MLIRLinalgTransforms
   MLIRMathToLibm
   MLIRMathToLLVM
   MLIRMemRefToLLVM
   MLIRPass
   MLIRReconcileUnrealizedCasts
   MLIRSCFToControlFlow
   MLIRSparseTensor
   MLIRSparseTensorTransforms
   MLIRTensorTransforms
   MLIRVectorToLLVM
   MLIRVectorTransforms
 )

mlir/lib/Dialect/SparseTensor/Pipelines/SparseTensorPipelines.cpp

void mlir::sparse_tensor::buildSparseCompiler(
   pm.addPass(func::createFuncBufferizePass());
   pm.addPass(arith::createConstantBufferizePass());
   pm.addNestedPass<func::FuncOp>(createTensorBufferizePass());
   pm.addNestedPass<func::FuncOp>(
       mlir::bufferization::createFinalizingBufferizePass());
   pm.addPass(createLowerAffinePass());
   pm.addPass(createConvertVectorToLLVMPass(options.lowerVectorToLLVMOptions()));
   pm.addPass(createMemRefToLLVMPass());
+  pm.addNestedPass<func::FuncOp>(createConvertComplexToStandardPass());
   pm.addNestedPass<func::FuncOp>(createConvertMathToLLVMPass());
+  pm.addPass(createConvertComplexToLLVMPass());
   pm.addPass(createConvertMathToLibmPass());
   pm.addPass(createConvertFuncToLLVMPass());
   pm.addPass(createReconcileUnrealizedCastsPass());
 }
 
 //===----------------------------------------------------------------------===//
 // Pipeline registration.
 //===----------------------------------------------------------------------===//

mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.cpp

PrimaryType mlir::sparse_tensor::primaryTypeEncoding(Type elemTp) {
   if (elemTp.isInteger(64))
     return PrimaryType::kI64;
   if (elemTp.isInteger(32))
     return PrimaryType::kI32;
   if (elemTp.isInteger(16))
     return PrimaryType::kI16;
   if (elemTp.isInteger(8))
     return PrimaryType::kI8;
+  if (auto complexTp = elemTp.dyn_cast<ComplexType>()) {
+    auto complexEltTp = complexTp.getElementType();
+    if (complexEltTp.isF64())
+      return PrimaryType::kC64;
+    if (complexEltTp.isF32())
+      return PrimaryType::kC32;
+  }
   llvm_unreachable("Unknown primary type");
 }
 
 StringRef mlir::sparse_tensor::primaryTypeFunctionSuffix(PrimaryType pt) {
   switch (pt) {
   case PrimaryType::kF64:
     return "F64";
   case PrimaryType::kF32:
     return "F32";
   case PrimaryType::kI64:
     return "I64";
   case PrimaryType::kI32:
     return "I32";
   case PrimaryType::kI16:
     return "I16";
   case PrimaryType::kI8:
     return "I8";
+  case PrimaryType::kC64:
+    return "C64";
+  case PrimaryType::kC32:
+    return "C32";
   }
   llvm_unreachable("Unknown PrimaryType");
 }
 
 StringRef mlir::sparse_tensor::primaryTypeFunctionSuffix(Type elemTp) {
   return primaryTypeFunctionSuffix(primaryTypeEncoding(elemTp));
 }
 

mlir/lib/Dialect/SparseTensor/Utils/CMakeLists.txt

 add_mlir_dialect_library(MLIRSparseTensorUtils
   Merger.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/SparseTensor
 
   LINK_LIBS PUBLIC
   MLIRArithmetic
+  MLIRComplex
   MLIRIR
   MLIRLinalg
 )

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

 //===- Merger.cpp - Implementation of iteration lattices ------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/SparseTensor/Utils/Merger.h"
 #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
+#include "mlir/Dialect/Complex/IR/Complex.h"
 #include "mlir/Dialect/Math/IR/Math.h"
 #include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
 
 #include "mlir/IR/Operation.h"
 #include "llvm/Support/Debug.h"
 
 namespace mlir {
 namespace sparse_tensor {
case kDivU:
     assert(!maybeZero(tensorExps[e].children.e1));
     return isSingleCondition(t, tensorExps[e].children.e0);
   case kShrS: // note: x >> inv only
   case kShrU:
   case kShlI:
     assert(isInvariant(tensorExps[e].children.e1));
     return isSingleCondition(t, tensorExps[e].children.e0);
   case kMulF:
+  case kMulC:
   case kMulI:
   case kAndI:
     if (isSingleCondition(t, tensorExps[e].children.e0))
       return isSingleCondition(t, tensorExps[e].children.e1) ||
              isInvariant(tensorExps[e].children.e1);
     if (isSingleCondition(t, tensorExps[e].children.e1))
       return isInvariant(tensorExps[e].children.e0);
     return false;
   case kAddF:
+  case kAddC:
   case kAddI:
     return isSingleCondition(t, tensorExps[e].children.e0) &&
            isSingleCondition(t, tensorExps[e].children.e1);
   default:
     return false;
   }
 }
 
static const char *kindToOpSymbol(Kind kind) {
   case kTruncI:
   case kBitCast:
     return "cast";
   case kBinaryBranch:
     return "binary_branch";
   case kUnary:
     return "unary";
   case kMulF:
-    return "*";
+  case kMulC:
   case kMulI:
     return "*";
   case kDivF:
-    return "/";
   case kDivS:
-    return "/";
   case kDivU:
     return "/";
   case kAddF:
-    return "+";
+  case kAddC:
   case kAddI:
     return "+";
   case kSubF:
-    return "-";
   case kSubI:
     return "-";
   case kAndI:
     return "&";
   case kOrI:
     return "|";
   case kXorI:
     return "^";
//      | absent() | present(y) |
       // invariant on the right.
       Block &absentBlock = absentRegion.front();
       YieldOp absentYield = cast<YieldOp>(absentBlock.getTerminator());
       Value absentVal = absentYield.result();
       unsigned rhs = addExp(kInvariant, absentVal);
       return takeDisj(kind, child0, buildLattices(rhs, i), unop);
     }
   case kMulF:
+  case kMulC:
   case kMulI:
   case kAndI:
     // A multiplicative operation only needs to be performed
     // for the conjunction of sparse iteration spaces.
     //
     //  x*y|!y | y |
     //  ---+---+---+
     //  !x | 0 | 0 |
     //   x | 0 |x*y|
+    //
+    // Note even here, 0*NaN=NaN and 0*Inf=NaN, but that is ignored.
     return takeConj(kind, // take binary conjunction
                     buildLattices(tensorExps[e].children.e0, i),
                     buildLattices(tensorExps[e].children.e1, i));
   case kDivF:
   case kDivS:
   case kDivU:
     // A division is tricky, since 0/0, 0/c, c/0 all have
     // specific outcomes for floating-point and integers.
     // Thus, we need to traverse the full iteration space.
     //
     //  x/y|!y | y |
     //  ---+---+---+
     //  !x |0/0|0/y|   FP: 0/0=NaN,c/0=Inf,0/c=0 with c true nonzero
     //   x |x/0|x/y|  INT: x/0=exception for any x
     //
     // TODO: for now we "fixed" this by only accepting x/c cases
     //       during expression building, so that the conjunction
     //       rules applies (viz. x/c = x*(1/c) as far as lattice
     //       construction is concerned).
     assert(!maybeZero(tensorExps[e].children.e1));
     return takeConj(kind, // take binary conjunction
                     buildLattices(tensorExps[e].children.e0, i),
                     buildLattices(tensorExps[e].children.e1, i));
   case kAddF:
+  case kAddC:
   case kAddI:
   case kSubF:
   case kSubI:
   case kOrI:
   case kXorI:
     // An additive operation needs to be performed
     // for the disjunction of sparse iteration spaces.
     //
Optional<unsigned> Merger::buildTensorExp(linalg::GenericOp op, Value v) {
   if (def->getNumOperands() == 2) {
     auto x = buildTensorExp(op, def->getOperand(0));
     auto y = buildTensorExp(op, def->getOperand(1));
     if (x.hasValue() && y.hasValue()) {
       unsigned e0 = x.getValue();
       unsigned e1 = y.getValue();
       if (isa<arith::MulFOp>(def))
         return addExp(kMulF, e0, e1);
+      if (isa<complex::MulOp>(def))
+        return addExp(kMulC, e0, e1);
       if (isa<arith::MulIOp>(def))
         return addExp(kMulI, e0, e1);
       if (isa<arith::DivFOp>(def) && !maybeZero(e1))
         return addExp(kDivF, e0, e1);
       if (isa<arith::DivSIOp>(def) && !maybeZero(e1))
         return addExp(kDivS, e0, e1);
       if (isa<arith::DivUIOp>(def) && !maybeZero(e1))
         return addExp(kDivU, e0, e1);
       if (isa<arith::AddFOp>(def))
         return addExp(kAddF, e0, e1);
+      if (isa<complex::AddOp>(def))
+        return addExp(kAddC, e0, e1);
       if (isa<arith::AddIOp>(def))
         return addExp(kAddI, e0, e1);
       if (isa<arith::SubFOp>(def))
         return addExp(kSubF, e0, e1);
       if (isa<arith::SubIOp>(def))
         return addExp(kSubI, e0, e1);
       if (isa<arith::AndIOp>(def))
         return addExp(kAndI, e0, e1);
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 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:
     return rewriter.create<arith::MulIOp>(loc, v0, v1);
   case kDivF:
     return rewriter.create<arith::DivFOp>(loc, v0, v1);
   case kDivS:
     return rewriter.create<arith::DivSIOp>(loc, v0, v1);
   case kDivU:
     return rewriter.create<arith::DivUIOp>(loc, v0, v1);
   case kAddF:
     return rewriter.create<arith::AddFOp>(loc, v0, v1);
+  case kAddC:
+    return rewriter.create<complex::AddOp>(loc, v0, v1);
   case kAddI:
     return rewriter.create<arith::AddIOp>(loc, v0, v1);
   case kSubF:
     return rewriter.create<arith::SubFOp>(loc, v0, v1);
   case kSubI:
     return rewriter.create<arith::SubIOp>(loc, v0, v1);
   case kAndI:
     return rewriter.create<arith::AndIOp>(loc, v0, v1);

mlir/lib/ExecutionEngine/SparseTensorUtils.cpp

 
 #include "mlir/ExecutionEngine/SparseTensorUtils.h"
 #include "mlir/ExecutionEngine/CRunnerUtils.h"
 
 #ifdef MLIR_CRUNNERUTILS_DEFINE_FUNCTIONS
 
 #include <algorithm>
 #include <cassert>
+#include <ccomplex>
 #include <cctype>
 #include <cinttypes>
 #include <cstdio>
 #include <cstdlib>
 #include <cstring>
 #include <fstream>
 #include <functional>
 #include <iostream>
 #include <limits>
 #include <numeric>
 #include <vector>
 
+using complex64 = std::complex<double>;
+using complex32 = std::complex<float>;
+
 //===----------------------------------------------------------------------===//
 //
 // Internal support for storing and reading sparse tensors.
 //
 // The following memory-resident sparse storage schemes are supported:
 //
 // (a) A coordinate scheme for temporarily storing and lexicographically
 //     sorting a sparse tensor by index (SparseTensorCOO).
public:
 
   /// Primary storage.
   virtual void getValues(std::vector<double> **) { fatal("valf64"); }
   virtual void getValues(std::vector<float> **) { fatal("valf32"); }
   virtual void getValues(std::vector<int64_t> **) { fatal("vali64"); }
   virtual void getValues(std::vector<int32_t> **) { fatal("vali32"); }
   virtual void getValues(std::vector<int16_t> **) { fatal("vali16"); }
   virtual void getValues(std::vector<int8_t> **) { fatal("vali8"); }
+  virtual void getValues(std::vector<complex64> **) { fatal("valc64"); }
+  virtual void getValues(std::vector<complex32> **) { fatal("valc32"); }
 
   /// Element-wise insertion in lexicographic index order.
   virtual void lexInsert(const uint64_t *, double) { fatal("insf64"); }
   virtual void lexInsert(const uint64_t *, float) { fatal("insf32"); }
   virtual void lexInsert(const uint64_t *, int64_t) { fatal("insi64"); }
   virtual void lexInsert(const uint64_t *, int32_t) { fatal("insi32"); }
   virtual void lexInsert(const uint64_t *, int16_t) { fatal("ins16"); }
   virtual void lexInsert(const uint64_t *, int8_t) { fatal("insi8"); }
+  virtual void lexInsert(const uint64_t *, complex64) { fatal("insc64"); }
+  virtual void lexInsert(const uint64_t *, complex32) { fatal("insc32"); }
 
   /// Expanded insertion.
   virtual void expInsert(uint64_t *, double *, bool *, uint64_t *, uint64_t) {
     fatal("expf64");
   }
   virtual void expInsert(uint64_t *, float *, bool *, uint64_t *, uint64_t) {
     fatal("expf32");
   }
   virtual void expInsert(uint64_t *, int64_t *, bool *, uint64_t *, uint64_t) {
     fatal("expi64");
   }
   virtual void expInsert(uint64_t *, int32_t *, bool *, uint64_t *, uint64_t) {
     fatal("expi32");
   }
   virtual void expInsert(uint64_t *, int16_t *, bool *, uint64_t *, uint64_t) {
     fatal("expi16");
   }
   virtual void expInsert(uint64_t *, int8_t *, bool *, uint64_t *, uint64_t) {
     fatal("expi8");
   }
+  virtual void expInsert(uint64_t *, complex64 *, bool *, uint64_t *,
+                         uint64_t) {
+    fatal("expc64");
+  }
+  virtual void expInsert(uint64_t *, complex32 *, bool *, uint64_t *,
+                         uint64_t) {
+    fatal("expc32");
+  }
 
   /// Finishes insertion.
   virtual void endInsert() = 0;
 
 protected:
   // Since this class is virtual, we must disallow public copying in
   // order to avoid "slicing".  Since this class has data members,
   // that means making copying protected.
static SparseTensorCOO<V> *openSparseTensorCOO(char *filename, uint64_t rank,
   // and the number of nonzeros as initial capacity.
   assert(rank == idata[0] && "rank mismatch");
   uint64_t nnz = idata[1];
   for (uint64_t r = 0; r < rank; r++)
     assert((shape[r] == 0 || shape[r] == idata[2 + r]) &&
            "dimension size mismatch");
   SparseTensorCOO<V> *tensor =
       SparseTensorCOO<V>::newSparseTensorCOO(rank, idata + 2, perm, nnz);
-  //  Read all nonzero elements.
+  // Read all nonzero elements.
   std::vector<uint64_t> indices(rank);
   for (uint64_t k = 0; k < nnz; k++) {
     if (!fgets(line, kColWidth, file)) {
       fprintf(stderr, "Cannot find next line of data in %s\n", filename);
       exit(1);
     }
     char *linePtr = line;
     for (uint64_t r = 0; r < rank; r++) {
#endif
   return tensor;
 }
 
 /// Converts a sparse tensor to an external COO-flavored format.
 template <typename V>
 static void fromMLIRSparseTensor(void *tensor, uint64_t *pRank, uint64_t *pNse,
                                  uint64_t **pShape, V **pValues,
                                  uint64_t **pIndices) {
+  assert(tensor);
   auto sparseTensor =
       static_cast<SparseTensorStorage<uint64_t, uint64_t, V> *>(tensor);
   uint64_t rank = sparseTensor->getRank();
   std::vector<uint64_t> perm(rank);
   std::iota(perm.begin(), perm.end(), 0);
   SparseTensorCOO<V> *coo = sparseTensor->toCOO(perm.data());
 
   const std::vector<Element<V>> &elements = coo->getElements();
_mlir_ciface_newSparseTensor(StridedMemRefType<DimLevelType, 1> *aref, // NOLINT
   CASE_SECSAME(OverheadType::kU32, PrimaryType::kI8, uint32_t, int8_t);
   CASE_SECSAME(OverheadType::kU16, PrimaryType::kI32, uint16_t, int32_t);
   CASE_SECSAME(OverheadType::kU16, PrimaryType::kI16, uint16_t, int16_t);
   CASE_SECSAME(OverheadType::kU16, PrimaryType::kI8, uint16_t, int8_t);
   CASE_SECSAME(OverheadType::kU8, PrimaryType::kI32, uint8_t, int32_t);
   CASE_SECSAME(OverheadType::kU8, PrimaryType::kI16, uint8_t, int16_t);
   CASE_SECSAME(OverheadType::kU8, PrimaryType::kI8, uint8_t, int8_t);
 
+  // Complex matrices with wide overhead.
+  CASE_SECSAME(OverheadType::kU64, PrimaryType::kC64, uint64_t, complex64);
+  CASE_SECSAME(OverheadType::kU64, PrimaryType::kC32, uint64_t, complex32);
+
   // Unsupported case (add above if needed).
   fputs("unsupported combination of types\n", stderr);
   exit(1);
 }
 
 /// Methods that provide direct access to pointers.
 IMPL_GETOVERHEAD(sparsePointers, index_type, getPointers)
 IMPL_GETOVERHEAD(sparsePointers64, uint64_t, getPointers)

 
 /// Methods that provide direct access to values.
 IMPL_SPARSEVALUES(sparseValuesF64, double, getValues)
 IMPL_SPARSEVALUES(sparseValuesF32, float, getValues)
 IMPL_SPARSEVALUES(sparseValuesI64, int64_t, getValues)
 IMPL_SPARSEVALUES(sparseValuesI32, int32_t, getValues)
 IMPL_SPARSEVALUES(sparseValuesI16, int16_t, getValues)
 IMPL_SPARSEVALUES(sparseValuesI8, int8_t, getValues)
+IMPL_SPARSEVALUES(sparseValuesC64, complex64, getValues)
+IMPL_SPARSEVALUES(sparseValuesC32, complex32, getValues)
 
 /// Helper to add value to coordinate scheme, one per value type.
 IMPL_ADDELT(addEltF64, double)
 IMPL_ADDELT(addEltF32, float)
 IMPL_ADDELT(addEltI64, int64_t)
 IMPL_ADDELT(addEltI32, int32_t)
 IMPL_ADDELT(addEltI16, int16_t)
 IMPL_ADDELT(addEltI8, int8_t)
+IMPL_ADDELT(addEltC64, complex64)
+IMPL_ADDELT(addEltC32ABI, complex32)
+// TODO: cleaner way to avoid ABI padding problem?
+void *_mlir_ciface_addEltC32(void *tensor, float r, float i,
+                             StridedMemRefType<index_type, 1> *iref,
+                             StridedMemRefType<index_type, 1> *pref) {
+  return _mlir_ciface_addEltC32ABI(tensor, complex32(r, i), iref, pref);
+}
 
 /// Helper to enumerate elements of coordinate scheme, one per value type.
 IMPL_GETNEXT(getNextF64, double)
 IMPL_GETNEXT(getNextF32, float)
 IMPL_GETNEXT(getNextI64, int64_t)
 IMPL_GETNEXT(getNextI32, int32_t)
 IMPL_GETNEXT(getNextI16, int16_t)
 IMPL_GETNEXT(getNextI8, int8_t)
+IMPL_GETNEXT(getNextC64, complex64)
+IMPL_GETNEXT(getNextC32, complex32)
 
 /// Insert elements in lexicographical index order, one per value type.
 IMPL_LEXINSERT(lexInsertF64, double)
 IMPL_LEXINSERT(lexInsertF32, float)
 IMPL_LEXINSERT(lexInsertI64, int64_t)
 IMPL_LEXINSERT(lexInsertI32, int32_t)
 IMPL_LEXINSERT(lexInsertI16, int16_t)
 IMPL_LEXINSERT(lexInsertI8, int8_t)
+IMPL_LEXINSERT(lexInsertC64, complex64)
+IMPL_LEXINSERT(lexInsertC32ABI, complex32)
+// TODO: cleaner way to avoid ABI padding problem?
+void _mlir_ciface_lexInsertC32(void *tensor,
+                               StridedMemRefType<index_type, 1> *cref, float r,
+                               float i) {
+  _mlir_ciface_lexInsertC32ABI(tensor, cref, complex32(r, i));
+}
 
 /// Insert using expansion, one per value type.
 IMPL_EXPINSERT(expInsertF64, double)
 IMPL_EXPINSERT(expInsertF32, float)
 IMPL_EXPINSERT(expInsertI64, int64_t)
 IMPL_EXPINSERT(expInsertI32, int32_t)
 IMPL_EXPINSERT(expInsertI16, int16_t)
 IMPL_EXPINSERT(expInsertI8, int8_t)
+IMPL_EXPINSERT(expInsertC64, complex64)
+IMPL_EXPINSERT(expInsertC32, complex32)
 
 #undef CASE
 #undef IMPL_SPARSEVALUES
 #undef IMPL_GETOVERHEAD
 #undef IMPL_ADDELT
 #undef IMPL_GETNEXT
 #undef IMPL_LEXINSERT
 #undef IMPL_EXPINSERT
void outSparseTensorI32(void *tensor, void *dest, bool sort) {
   return outSparseTensor<int32_t>(tensor, dest, sort);
 }
 void outSparseTensorI16(void *tensor, void *dest, bool sort) {
   return outSparseTensor<int16_t>(tensor, dest, sort);
 }
 void outSparseTensorI8(void *tensor, void *dest, bool sort) {
   return outSparseTensor<int8_t>(tensor, dest, sort);
 }
+void outSparseTensorC64(void *tensor, void *dest, bool sort) {
+  return outSparseTensor<complex64>(tensor, dest, sort);
+}
+void outSparseTensorC32(void *tensor, void *dest, bool sort) {
+  return outSparseTensor<complex32>(tensor, dest, sort);
+}
 
 //===----------------------------------------------------------------------===//
 //
 // Public API with methods that accept C-style data structures to interact
 // with sparse tensors, which are only visible as opaque pointers externally.
 // These methods can be used both by MLIR compiler-generated code as well as by
 // an external runtime that wants to interact with MLIR compiler-generated code.
 //
void delSparseTensorCOO##VNAME(void *coo) {                                  \
     delete static_cast<SparseTensorCOO<V> *>(coo);                             \
   }
 IMPL_DELCOO(F64, double)
 IMPL_DELCOO(F32, float)
 IMPL_DELCOO(I64, int64_t)
 IMPL_DELCOO(I32, int32_t)
 IMPL_DELCOO(I16, int16_t)
 IMPL_DELCOO(I8, int8_t)
+IMPL_DELCOO(C64, complex64)
+IMPL_DELCOO(C32, complex32)
 #undef IMPL_DELCOO
 
 /// Initializes sparse tensor from a COO-flavored format expressed using C-style
 /// data structures. The expected parameters are:
 ///
 ///   rank:    rank of tensor
 ///   nse:     number of specified elements (usually the nonzeros)
 ///   shape:   array with dimension size for each rank
return toMLIRSparseTensor<int16_t>(rank, nse, shape, values, indices, perm,
                                      sparse);
 }
 void *convertToMLIRSparseTensorI8(uint64_t rank, uint64_t nse, uint64_t *shape,
                                   int8_t *values, uint64_t *indices,
                                   uint64_t *perm, uint8_t *sparse) {
   return toMLIRSparseTensor<int8_t>(rank, nse, shape, values, indices, perm,
                                     sparse);
 }
+void *convertToMLIRSparseTensorC64(uint64_t rank, uint64_t nse, uint64_t *shape,
+                                   complex64 *values, uint64_t *indices,
+                                   uint64_t *perm, uint8_t *sparse) {
+  return toMLIRSparseTensor<complex64>(rank, nse, shape, values, indices, perm,
+                                       sparse);
+}
+void *convertToMLIRSparseTensorC32(uint64_t rank, uint64_t nse, uint64_t *shape,
+                                   complex32 *values, uint64_t *indices,
+                                   uint64_t *perm, uint8_t *sparse) {
+  return toMLIRSparseTensor<complex32>(rank, nse, shape, values, indices, perm,
+                                       sparse);
+}
 
 /// Converts a sparse tensor to COO-flavored format expressed using C-style
 /// data structures. The expected output parameters are pointers for these
 /// values:
 ///
 ///   rank:    rank of tensor
 ///   nse:     number of specified elements (usually the nonzeros)
 ///   shape:   array with dimension size for each rank
void convertFromMLIRSparseTensorI16(void *tensor, uint64_t *pRank,
                                     int16_t **pValues, uint64_t **pIndices) {
   fromMLIRSparseTensor<int16_t>(tensor, pRank, pNse, pShape, pValues, pIndices);
 }
 void convertFromMLIRSparseTensorI8(void *tensor, uint64_t *pRank,
                                    uint64_t *pNse, uint64_t **pShape,
                                    int8_t **pValues, uint64_t **pIndices) {
   fromMLIRSparseTensor<int8_t>(tensor, pRank, pNse, pShape, pValues, pIndices);
 }
+void convertFromMLIRSparseTensorC64(void *tensor, uint64_t *pRank,
+                                    uint64_t *pNse, uint64_t **pShape,
+                                    complex64 **pValues, uint64_t **pIndices) {
+  fromMLIRSparseTensor<complex64>(tensor, pRank, pNse, pShape, pValues,
+                                  pIndices);
+}
+void convertFromMLIRSparseTensorC32(void *tensor, uint64_t *pRank,
+                                    uint64_t *pNse, uint64_t **pShape,
+                                    complex32 **pValues, uint64_t **pIndices) {
+  fromMLIRSparseTensor<complex32>(tensor, pRank, pNse, pShape, pValues,
+                                  pIndices);
+}
 
 } // extern "C"
 
 #endif // MLIR_CRUNNERUTILS_DEFINE_FUNCTIONS

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_complex32.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)>,  // b (in)
+    affine_map<(i) -> (i)>   // x (out)
+  ],
+  iterator_types = ["parallel"],
+  doc = "x(i) = a(i) OP b(i)"
+}
+
+module {
+  func.func @cadd(%arga: tensor<?xcomplex<f32>, #SparseVector>,
+                  %argb: tensor<?xcomplex<f32>, #SparseVector>)
+                      -> tensor<?xcomplex<f32>, #SparseVector> {
+    %c = arith.constant 0 : index
+    %d = tensor.dim %arga, %c : tensor<?xcomplex<f32>, #SparseVector>
+    %xv = sparse_tensor.init [%d] : tensor<?xcomplex<f32>, #SparseVector>
+    %0 = linalg.generic #trait_op
+       ins(%arga, %argb: tensor<?xcomplex<f32>, #SparseVector>,
+                         tensor<?xcomplex<f32>, #SparseVector>)
+        outs(%xv: tensor<?xcomplex<f32>, #SparseVector>) {
+        ^bb(%a: complex<f32>, %b: complex<f32>, %x: complex<f32>):
+          %1 = complex.add %a, %b : complex<f32>
+          linalg.yield %1 : complex<f32>
+    } -> tensor<?xcomplex<f32>, #SparseVector>
+    return %0 : tensor<?xcomplex<f32>, #SparseVector>
+  }
+
+  func.func @cmul(%arga: tensor<?xcomplex<f32>, #SparseVector>,
+                  %argb: tensor<?xcomplex<f32>, #SparseVector>)
+                      -> tensor<?xcomplex<f32>, #SparseVector> {
+    %c = arith.constant 0 : index
+    %d = tensor.dim %arga, %c : tensor<?xcomplex<f32>, #SparseVector>
+    %xv = sparse_tensor.init [%d] : tensor<?xcomplex<f32>, #SparseVector>
+    %0 = linalg.generic #trait_op
+       ins(%arga, %argb: tensor<?xcomplex<f32>, #SparseVector>,
+                         tensor<?xcomplex<f32>, #SparseVector>)
+        outs(%xv: tensor<?xcomplex<f32>, #SparseVector>) {
+        ^bb(%a: complex<f32>, %b: complex<f32>, %x: complex<f32>):
+          %1 = complex.mul %a, %b : complex<f32>
+          linalg.yield %1 : complex<f32>
+    } -> tensor<?xcomplex<f32>, #SparseVector>
+    return %0 : tensor<?xcomplex<f32>, #SparseVector>
+  }
+
+  func.func @dump(%arg0: tensor<?xcomplex<f32>, #SparseVector>, %d: index) {
+    %c0 = arith.constant 0 : index
+    %c1 = arith.constant 1 : index
+    %mem = sparse_tensor.values %arg0 : tensor<?xcomplex<f32>, #SparseVector> to memref<?xcomplex<f32>>
+    scf.for %i = %c0 to %d step %c1 {
+       %v = memref.load %mem[%i] : memref<?xcomplex<f32>>
+       %real = complex.re %v : complex<f32>
+       %imag = complex.im %v : complex<f32>
+       vector.print %real : f32
+       vector.print %imag : f32
+    }
+    return
+  }
+
+  // Driver method to call and verify complex kernels.
+  func.func @entry() {
+    // Setup sparse vectors.
+    %v1 = arith.constant sparse<
+       [ [0], [28], [31] ],
+         [ (511.13, 2.0), (3.0, 4.0), (5.0, 6.0) ] > : tensor<32xcomplex<f32>>
+    %v2 = arith.constant sparse<
+       [ [1], [28], [31] ],
+         [ (1.0, 0.0), (2.0, 0.0), (3.0, 0.0) ] > : tensor<32xcomplex<f32>>
+    %sv1 = sparse_tensor.convert %v1 : tensor<32xcomplex<f32>> to tensor<?xcomplex<f32>, #SparseVector>
+    %sv2 = sparse_tensor.convert %v2 : tensor<32xcomplex<f32>> to tensor<?xcomplex<f32>, #SparseVector>
+
+    // Call sparse vector kernels.
+    %0 = call @cadd(%sv1, %sv2)
+       : (tensor<?xcomplex<f32>, #SparseVector>,
+          tensor<?xcomplex<f32>, #SparseVector>) -> tensor<?xcomplex<f32>, #SparseVector>
+    %1 = call @cmul(%sv1, %sv2)
+       : (tensor<?xcomplex<f32>, #SparseVector>,
+          tensor<?xcomplex<f32>, #SparseVector>) -> tensor<?xcomplex<f32>, #SparseVector>
+
+    //
+    // Verify the results.
+    //
+    // CHECK: 511.13
+    // CHECK-NEXT: 2
+    // CHECK-NEXT: 1
+    // CHECK-NEXT: 0
+    // CHECK-NEXT: 5
+    // CHECK-NEXT: 4
+    // CHECK-NEXT: 8
+    // CHECK-NEXT: 6
+    // CHECK-NEXT: 6
+    // CHECK-NEXT: 8
+    // CHECK-NEXT: 15
+    // CHECK-NEXT: 18
+    //
+    %d1 = arith.constant 4 : index
+    %d2 = arith.constant 2 : index
+    call @dump(%0, %d1) : (tensor<?xcomplex<f32>, #SparseVector>, index) -> ()
+    call @dump(%1, %d2) : (tensor<?xcomplex<f32>, #SparseVector>, index) -> ()
+
+    // Release the resources.
+    sparse_tensor.release %sv1 : tensor<?xcomplex<f32>, #SparseVector>
+    sparse_tensor.release %sv2 : tensor<?xcomplex<f32>, #SparseVector>
+    sparse_tensor.release %0 : tensor<?xcomplex<f32>, #SparseVector>
+    sparse_tensor.release %1 : tensor<?xcomplex<f32>, #SparseVector>
+    return
+  }
+}

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_complex64.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)>,  // b (in)
+    affine_map<(i) -> (i)>   // x (out)
+  ],
+  iterator_types = ["parallel"],
+  doc = "x(i) = a(i) OP b(i)"
+}
+
+module {
+  func.func @cadd(%arga: tensor<?xcomplex<f64>, #SparseVector>,
+                  %argb: tensor<?xcomplex<f64>, #SparseVector>)
+                      -> tensor<?xcomplex<f64>, #SparseVector> {
+    %c = arith.constant 0 : index
+    %d = tensor.dim %arga, %c : tensor<?xcomplex<f64>, #SparseVector>
+    %xv = sparse_tensor.init [%d] : tensor<?xcomplex<f64>, #SparseVector>
+    %0 = linalg.generic #trait_op
+       ins(%arga, %argb: tensor<?xcomplex<f64>, #SparseVector>,
+                         tensor<?xcomplex<f64>, #SparseVector>)
+        outs(%xv: tensor<?xcomplex<f64>, #SparseVector>) {
+        ^bb(%a: complex<f64>, %b: complex<f64>, %x: complex<f64>):
+          %1 = complex.add %a, %b : complex<f64>
+          linalg.yield %1 : complex<f64>
+    } -> tensor<?xcomplex<f64>, #SparseVector>
+    return %0 : tensor<?xcomplex<f64>, #SparseVector>
+  }
+
+  func.func @cmul(%arga: tensor<?xcomplex<f64>, #SparseVector>,
+                  %argb: tensor<?xcomplex<f64>, #SparseVector>)
+                      -> tensor<?xcomplex<f64>, #SparseVector> {
+    %c = arith.constant 0 : index
+    %d = tensor.dim %arga, %c : tensor<?xcomplex<f64>, #SparseVector>
+    %xv = sparse_tensor.init [%d] : tensor<?xcomplex<f64>, #SparseVector>
+    %0 = linalg.generic #trait_op
+       ins(%arga, %argb: tensor<?xcomplex<f64>, #SparseVector>,
+                         tensor<?xcomplex<f64>, #SparseVector>)
+        outs(%xv: tensor<?xcomplex<f64>, #SparseVector>) {
+        ^bb(%a: complex<f64>, %b: complex<f64>, %x: complex<f64>):
+          %1 = complex.mul %a, %b : complex<f64>
+          linalg.yield %1 : complex<f64>
+    } -> tensor<?xcomplex<f64>, #SparseVector>
+    return %0 : tensor<?xcomplex<f64>, #SparseVector>
+  }
+
+  func.func @dump(%arg0: tensor<?xcomplex<f64>, #SparseVector>, %d: index) {
+    %c0 = arith.constant 0 : index
+    %c1 = arith.constant 1 : index
+    %mem = sparse_tensor.values %arg0 : tensor<?xcomplex<f64>, #SparseVector> to memref<?xcomplex<f64>>
+    scf.for %i = %c0 to %d step %c1 {
+       %v = memref.load %mem[%i] : memref<?xcomplex<f64>>
+       %real = complex.re %v : complex<f64>
+       %imag = complex.im %v : complex<f64>
+       vector.print %real : f64
+       vector.print %imag : f64
+    }
+    return
+  }
+
+  // Driver method to call and verify complex kernels.
+  func.func @entry() {
+    // Setup sparse vectors.
+    %v1 = arith.constant sparse<
+       [ [0], [28], [31] ],
+         [ (511.13, 2.0), (3.0, 4.0), (5.0, 6.0) ] > : tensor<32xcomplex<f64>>
+    %v2 = arith.constant sparse<
+       [ [1], [28], [31] ],
+         [ (1.0, 0.0), (2.0, 0.0), (3.0, 0.0) ] > : tensor<32xcomplex<f64>>
+    %sv1 = sparse_tensor.convert %v1 : tensor<32xcomplex<f64>> to tensor<?xcomplex<f64>, #SparseVector>
+    %sv2 = sparse_tensor.convert %v2 : tensor<32xcomplex<f64>> to tensor<?xcomplex<f64>, #SparseVector>
+
+    // Call sparse vector kernels.
+    %0 = call @cadd(%sv1, %sv2)
+       : (tensor<?xcomplex<f64>, #SparseVector>,
+          tensor<?xcomplex<f64>, #SparseVector>) -> tensor<?xcomplex<f64>, #SparseVector>
+    %1 = call @cmul(%sv1, %sv2)
+       : (tensor<?xcomplex<f64>, #SparseVector>,
+          tensor<?xcomplex<f64>, #SparseVector>) -> tensor<?xcomplex<f64>, #SparseVector>
+
+    //
+    // Verify the results.
+    //
+    // CHECK: 511.13
+    // CHECK-NEXT: 2
+    // CHECK-NEXT: 1
+    // CHECK-NEXT: 0
+    // CHECK-NEXT: 5
+    // CHECK-NEXT: 4
+    // CHECK-NEXT: 8
+    // CHECK-NEXT: 6
+    // CHECK-NEXT: 6
+    // CHECK-NEXT: 8
+    // CHECK-NEXT: 15
+    // CHECK-NEXT: 18
+    //
+    %d1 = arith.constant 4 : index
+    %d2 = arith.constant 2 : index
+    call @dump(%0, %d1) : (tensor<?xcomplex<f64>, #SparseVector>, index) -> ()
+    call @dump(%1, %d2) : (tensor<?xcomplex<f64>, #SparseVector>, index) -> ()
+
+    // 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>
+    return
+  }
+}

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

 
 cc_library(
     name = "SparseTensorUtils",
     srcs = glob(["lib/Dialect/SparseTensor/Utils/*.cpp"]),
     hdrs = glob(["include/mlir/Dialect/SparseTensor/Utils/*.h"]),
     includes = ["include"],
     deps = [
         ":ArithmeticDialect",
+        ":ComplexDialect",
         ":IR",
         ":LinalgOps",
         ":MathDialect",
         ":SparseTensor",
         "//llvm:Support",
     ],
 )