[mlir][sparse] Add complex number reading from files.

[bixia]

Summary

Support complex numbers for Matrix Market Exchange Formats. Add a test case.

[aartbik]

mlir/lib/ExecutionEngine/SparseTensorUtils.cpp

1264–1265

NOTE: this now should distinguish between real, complex, and probably even integer!

1360–1370

note that for the MM in particular, we should also verify that the data is marked real/complex/integer going forward, and verify that this is compatible with the actual data type <V> (assert, or otherwise fail when imcompatible)

1372

I would add a comment on reading two doubles to make a complex here first, just so that it stays similar to L1317-1319. Otherwise, directly addressing the pattern issue reads a bit too terse.

1386–1387

This is of course the right way forward, but did you measure perf impact of introducing more calls along the path that is executed for each nonzero? Perhaps try reading in one of the larger tensors before/after this change. Would marking methods as inline help here?

mlir/test/Integration/data/test_symmetric_complex.mtx

2

This must be "complex", not "real' following proper MM format

[bixia]

Add data type check to MME header parsing.
Make the functions to read a COO value inline functions.

[bixia]

mlir/lib/ExecutionEngine/SparseTensorUtils.cpp

1360–1370

Make the functions inline functions.

[bixia]

Fix a typo.

[aartbik]

mlir/lib/ExecutionEngine/SparseTensorUtils.cpp

1237

This is of course a reasonable way to do it. But one may a bit surprised by having the fatal assert while reading the header.

I feel that simply reading and storing the field in the SparseTensorFile class as an enum (with pattern, integer, real, complex, undef {to deal with FROSTT}) and then having getters will let the client decide what to do once the header is properly parsed seemed a bit more along the direction the refactoring in a class was taking.

I don't feel super strongly about it, so this is more a WDYT? comment

[aartbik]

mlir/lib/ExecutionEngine/SparseTensorUtils.cpp

1263

this is a very narrow subset of possibilities (i.e. must match exactly).

I can see cases where we allow ints to be assigned to real and vv (as we did before),
or real/integer to complex with imaginary part zero.

All such cases seem a lot easier to decide in the "client" (even if that is just us right now ;-)
with the enum approach

[bixia]

Use enum to store the value type to SparseTensorFile and check type compatible in openSparseTensorCOO.

[bixia]

mlir/lib/ExecutionEngine/SparseTensorUtils.cpp

1237

Add enum ValueKind to store the value type information and use the information in the reader.

1360–1370

Move this check to the reader.

[aartbik]

I love it! Thanks for making the change. Two last nits, but good to go after that.

mlir/lib/ExecutionEngine/SparseTensorUtils.cpp

1299–1349

Maybe add a comment, e.g.

// The FROSTT format does not define the data type of the nonzero elements.
valueKind_ = ValueKind::kUndefined;

1367

This does not seem to be exhaustive, i.e. complex vs isint seems an error too

given that we have accepted real and int before, I would be perfectly okay that the only case we reject right now is complex file vs. non-complex data (which we can later refine into zero-imaginary part)

[bixia]

Add a comment.

mlir/lib/ExecutionEngine/SparseTensorUtils.cpp

1367

complex data vs integer tensor is an error. Maybe you misread tensorIsReal? It also includes integer.

[bixia]

There is one place where isValid() should be used, not isValid.

mlir/lib/ExecutionEngine/SparseTensorUtils.cpp

for (char *c = token; *c; c++)
     *c = tolower(*c);
   return token;
 }
 
 /// This class abstracts over the information stored in file headers,
 /// as well as providing the buffers and methods for parsing those headers.
 class SparseTensorFile final {
 public:
+  enum class ValueKind {
+    kInvalid = 0,
+    kPattern = 1,
+    kReal = 2,
+    kInteger = 3,
+    kComplex = 4,
+    kUndefined = 5
+  };
+
   explicit SparseTensorFile(char *filename) : filename(filename) {
     assert(filename && "Received nullptr for filename");
   }
 
   // Disallows copying, to avoid duplicating the `file` pointer.
   SparseTensorFile(const SparseTensorFile &) = delete;
   SparseTensorFile &operator=(const SparseTensorFile &) = delete;
 
public:
   void readHeader() {
     assert(file && "Attempt to readHeader() before openFile()");
     if (strstr(filename, ".mtx"))
       readMMEHeader();
     else if (strstr(filename, ".tns"))
       readExtFROSTTHeader();
     else
       FATAL("Unknown format %s\n", filename);
-    assert(isValid && "Failed to read the header");
+    assert(isValid() && "Failed to read the header");
   }
 
+  ValueKind getValueKind() const { return valueKind_; }
+
+  bool isValid() const { return valueKind_ != ValueKind::kUndefined; }
+
   /// Gets the MME "pattern" property setting.  Is only valid after
   /// parsing the header.
   bool isPattern() const {
-    assert(isValid && "Attempt to isPattern() before readHeader()");
-    return isPattern_;
+    assert(isValid() && "Attempt to isPattern() before readHeader()");
+    return valueKind_ == ValueKind::kPattern;
   }
 
   /// Gets the MME "symmetric" property setting.  Is only valid after
   /// parsing the header.
   bool isSymmetric() const {
-    assert(isValid && "Attempt to isSymmetric() before readHeader()");
+    assert(isValid() && "Attempt to isSymmetric() before readHeader()");
     return isSymmetric_;
   }
 
   /// Gets the rank of the tensor.  Is only valid after parsing the header.
   uint64_t getRank() const {
-    assert(isValid && "Attempt to getRank() before readHeader()");
+    assert(isValid() && "Attempt to getRank() before readHeader()");
     return idata[0];
   }
 
   /// Gets the number of non-zeros.  Is only valid after parsing the header.
   uint64_t getNNZ() const {
-    assert(isValid && "Attempt to getNNZ() before readHeader()");
+    assert(isValid() && "Attempt to getNNZ() before readHeader()");
     return idata[1];
   }
 
   /// Gets the dimension-sizes array.  The pointer itself is always
   /// valid; however, the values stored therein are only valid after
   /// parsing the header.
   const uint64_t *getDimSizes() const { return idata + 2; }
 
public:
   }
 
 private:
   void readMMEHeader();
   void readExtFROSTTHeader();
 
   const char *filename;
   FILE *file = nullptr;
-  bool isValid = false;
-  bool isPattern_ = false;
+  ValueKind valueKind_ = ValueKind::kInvalid;
   bool isSymmetric_ = false;
   uint64_t idata[512];
   char line[kColWidth];
 };
 
 /// Read the MME header of a general sparse matrix of type real.
 void SparseTensorFile::readMMEHeader() {
   char header[64];
   char object[64];
   char format[64];
   char field[64];
   char symmetry[64];
   // Read header line.
   if (fscanf(file, "%63s %63s %63s %63s %63s\n", header, object, format, field,
              symmetry) != 5)
     FATAL("Corrupt header in %s\n", filename);
-  // Set properties
-  isPattern_ = (strcmp(toLower(field), "pattern") == 0);
+  // Process `field`, which specify pattern or the data type of the values.
+  if (strcmp(toLower(field), "pattern") == 0)
+    valueKind_ = ValueKind::kPattern;
+  else if (strcmp(toLower(field), "real") == 0)
+    valueKind_ = ValueKind::kReal;
+  else if (strcmp(toLower(field), "integer") == 0)
+    valueKind_ = ValueKind::kInteger;
+  else if (strcmp(toLower(field), "complex") == 0)
+    valueKind_ = ValueKind::kComplex;
+  else
+    FATAL("Unexpected header field value in %s\n", filename);
+
+  // Set properties.
   isSymmetric_ = (strcmp(toLower(symmetry), "symmetric") == 0);
   // Make sure this is a general sparse matrix.
   if (strcmp(toLower(header), "%%matrixmarket") ||
       strcmp(toLower(object), "matrix") ||
       strcmp(toLower(format), "coordinate") ||
-      (strcmp(toLower(field), "real") && !isPattern_) ||
       (strcmp(toLower(symmetry), "general") && !isSymmetric_))
     FATAL("Cannot find a general sparse matrix in %s\n", filename);
   // Skip comments.
   while (true) {
     readLine();
     if (line[0] != '%')
       break;
   }
   // Next line contains M N NNZ.
   idata[0] = 2; // rank
   if (sscanf(line, "%" PRIu64 "%" PRIu64 "%" PRIu64 "\n", idata + 2, idata + 3,
              idata + 1) != 3)
     FATAL("Cannot find size in %s\n", filename);
-  isValid = true;
 }
 
 /// Read the "extended" FROSTT header. Although not part of the documented
 /// format, we assume that the file starts with optional comments followed
 /// by two lines that define the rank, the number of nonzeros, and the
 /// dimensions sizes (one per rank) of the sparse tensor.
 void SparseTensorFile::readExtFROSTTHeader() {
   // Skip comments.
   while (true) {
     readLine();
     if (line[0] != '#')
       break;
   }
   // Next line contains RANK and NNZ.
   if (sscanf(line, "%" PRIu64 "%" PRIu64 "\n", idata, idata + 1) != 2)
     FATAL("Cannot find metadata in %s\n", filename);
   // Followed by a line with the dimension sizes (one per rank).
   for (uint64_t r = 0; r < idata[0]; r++)
     if (fscanf(file, "%" PRIu64, idata + 2 + r) != 1)
       FATAL("Cannot find dimension size %s\n", filename);
   readLine(); // end of line
-  isValid = true;
+  // The FROSTT format does not define the data type of the nonzero elements.
+  valueKind_ = ValueKind::kUndefined;
+}
+
+// Adds a value to a tensor in coordinate scheme. If is_symmetric_value is true,
+// also adds the value to its symmetric location.
+template <typename T, typename V>
+static inline void addValue(T *coo, V value,
+                            const std::vector<uint64_t> indices,
+                            bool is_symmetric_value) {
+  // TODO: <https://github.com/llvm/llvm-project/issues/54179>
+  coo->add(indices, value);
+  // We currently chose to deal with symmetric matrices by fully constructing
+  // them. In the future, we may want to make symmetry implicit for storage
+  // reasons.
+  if (is_symmetric_value)
+    coo->add({indices[1], indices[0]}, value);
+}
+
+// Reads an element of a complex type for the current indices in coordinate
+// scheme.
+template <typename V>
+static inline void readCOOValue(SparseTensorCOO<std::complex<V>> *coo,
+                                const std::vector<uint64_t> indices,
+                                char **linePtr, bool is_pattern,
+                                bool add_symmetric_value) {
+  // Read two values to make a complex. The external formats always store
+  // numerical values with the type double, but we cast these values to the
+  // sparse tensor object type. For a pattern tensor, we arbitrarily pick the
+  // value 1 for all entries.
+  V re = is_pattern ? 1.0 : strtod(*linePtr, linePtr);
+  V im = is_pattern ? 1.0 : strtod(*linePtr, linePtr);
+  std::complex<V> value = {re, im};
+  addValue(coo, value, indices, add_symmetric_value);
+}
+
+// Reads an element of a non-complex type for the current indices in coordinate
+// scheme.
+template <typename V,
+          typename std::enable_if<
+              !std::is_same<std::complex<float>, V>::value &&
+              !std::is_same<std::complex<double>, V>::value>::type * = nullptr>
+static void inline readCOOValue(SparseTensorCOO<V> *coo,
+                                const std::vector<uint64_t> indices,
+                                char **linePtr, bool is_pattern,
+                                bool is_symmetric_value) {
+  // The external formats always store these numerical values with the type
+  // double, but we cast these values to the sparse tensor object type.
+  // For a pattern tensor, we arbitrarily pick the value 1 for all entries.
+  double value = is_pattern ? 1.0 : strtod(*linePtr, linePtr);
+  addValue(coo, value, indices, is_symmetric_value);
 }
 
 /// Reads a sparse tensor with the given filename into a memory-resident
 /// sparse tensor in coordinate scheme.
 template <typename V>
-static SparseTensorCOO<V> *openSparseTensorCOO(char *filename, uint64_t rank,
-                                               const uint64_t *shape,
-                                               const uint64_t *perm) {
+static SparseTensorCOO<V> *
+openSparseTensorCOO(char *filename, uint64_t rank, const uint64_t *shape,
+                    const uint64_t *perm, PrimaryType valTp) {
   SparseTensorFile stfile(filename);
   stfile.openFile();
   stfile.readHeader();
+  // Check tensor element type against the value type in the input file.
+  SparseTensorFile::ValueKind valueKind = stfile.getValueKind();
+  bool tensorIsInteger =
+      (valTp >= PrimaryType::kI64 && valTp <= PrimaryType::kI8);
+  bool tensorIsReal = (valTp >= PrimaryType::kF64 && valTp <= PrimaryType::kI8);
+  if ((valueKind == SparseTensorFile::ValueKind::kReal && tensorIsInteger) ||
+      (valueKind == SparseTensorFile::ValueKind::kComplex && tensorIsReal)) {
+    FATAL("Tensor element type %d not compatible with values in file %s\n",
+          valTp, filename);
+  }
   stfile.assertMatchesShape(rank, shape);
   // Prepare sparse tensor object with per-dimension sizes
   // and the number of nonzeros as initial capacity.
   uint64_t nnz = stfile.getNNZ();
   auto *coo = SparseTensorCOO<V>::newSparseTensorCOO(rank, stfile.getDimSizes(),
                                                      perm, nnz);
   // Read all nonzero elements.
   std::vector<uint64_t> indices(rank);
   for (uint64_t k = 0; k < nnz; k++) {
     char *linePtr = stfile.readLine();
     for (uint64_t r = 0; r < rank; r++) {
       uint64_t idx = strtoul(linePtr, &linePtr, 10);
       // Add 0-based index.
       indices[perm[r]] = idx - 1;
     }
-    // The external formats always store the numerical values with the type
-    // double, but we cast these values to the sparse tensor object type.
-    // For a pattern tensor, we arbitrarily pick the value 1 for all entries.
-    double value = stfile.isPattern() ? 1.0 : strtod(linePtr, &linePtr);
-    // TODO: <https://github.com/llvm/llvm-project/issues/54179>
-    coo->add(indices, value);
-    // We currently chose to deal with symmetric matrices by fully constructing
-    // them. In the future, we may want to make symmetry implicit for storage
-    // reasons.
-    if (stfile.isSymmetric() && indices[0] != indices[1])
-      coo->add({indices[1], indices[0]}, value);
+    readCOOValue(coo, indices, &linePtr, stfile.isPattern(),
+                 stfile.isSymmetric() && indices[0] != indices[1]);
   }
   // Close the file and return tensor.
   stfile.closeFile();
   return coo;
 }
 
 /// Writes the sparse tensor to `dest` in extended FROSTT format.
 template <typename V>

 //===----------------------------------------------------------------------===//
 
 #define CASE(p, i, v, P, I, V)                                                 \
   if (ptrTp == (p) && indTp == (i) && valTp == (v)) {                          \
     SparseTensorCOO<V> *coo = nullptr;                                         \
     if (action <= Action::kFromCOO) {                                          \
       if (action == Action::kFromFile) {                                       \
         char *filename = static_cast<char *>(ptr);                             \
-        coo = openSparseTensorCOO<V>(filename, rank, shape, perm);             \
+        coo = openSparseTensorCOO<V>(filename, rank, shape, perm, v);          \
       } else if (action == Action::kFromCOO) {                                 \
         coo = static_cast<SparseTensorCOO<V> *>(ptr);                          \
       } else {                                                                 \
         assert(action == Action::kEmpty);                                      \
       }                                                                        \
       auto *tensor = SparseTensorStorage<P, I, V>::newSparseTensor(            \
           rank, shape, perm, sparsity, coo);                                   \
       if (action == Action::kFromFile)                                         \

mlir/test/CMakeLists.txt

if (MLIR_INCLUDE_INTEGRATION_TESTS)
   if(NOT TARGET ${LLVM_NATIVE_ARCH})
     message(FATAL_ERROR "MLIR_INCLUDE_INTEGRATION_TESTS requires a native target")
   endif()
 
   # Copy test data over.
   file(COPY ${CMAKE_CURRENT_SOURCE_DIR}/Integration/data/mttkrp_b.tns
             ${CMAKE_CURRENT_SOURCE_DIR}/Integration/data/test.mtx
             ${CMAKE_CURRENT_SOURCE_DIR}/Integration/data/test_symmetric.mtx
+            ${CMAKE_CURRENT_SOURCE_DIR}/Integration/data/test_symmetric_complex.mtx
             ${CMAKE_CURRENT_SOURCE_DIR}/Integration/data/test.tns
             ${CMAKE_CURRENT_SOURCE_DIR}/Integration/data/wide.mtx
           DESTINATION ${MLIR_INTEGRATION_TEST_DIR}/data/)
 endif()
 
 llvm_canonicalize_cmake_booleans(
   LLVM_BUILD_EXAMPLES
   MLIR_ENABLE_BINDINGS_PYTHON

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

+// RUN: mlir-opt %s --sparse-compiler | \
+// RUN: TENSOR0="%mlir_integration_test_dir/data/test_symmetric_complex.mtx" \
+// 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
+
+!Filename = !llvm.ptr<i8>
+
+#SparseMatrix = #sparse_tensor.encoding<{
+  dimLevelType = [ "compressed", "compressed" ]
+}>
+
+#trait_sum_reduce = {
+  indexing_maps = [
+    affine_map<(i,j) -> (i,j)>, // A
+    affine_map<(i,j) -> ()>     // x (out)
+  ],
+  iterator_types = ["reduction", "reduction"],
+  doc = "x += A(i,j)"
+}
+
+//
+// Integration test that lowers a kernel annotated as sparse to
+// actual sparse code, initializes a matching sparse storage scheme
+// from file, and runs the resulting code with the JIT compiler.
+//
+module {
+  //
+  // A kernel that sum-reduces a matrix to a single scalar.
+  //
+  func.func @kernel_sum_reduce(%arga: tensor<?x?xcomplex<f64>, #SparseMatrix>,
+                          %argx: tensor<complex<f64>> {linalg.inplaceable = true}) -> tensor<complex<f64>> {
+    %0 = linalg.generic #trait_sum_reduce
+      ins(%arga: tensor<?x?xcomplex<f64>, #SparseMatrix>)
+      outs(%argx: tensor<complex<f64>>) {
+      ^bb(%a: complex<f64>, %x: complex<f64>):
+        %0 = complex.add %x, %a : complex<f64>
+        linalg.yield %0 : complex<f64>
+    } -> tensor<complex<f64>>
+    return %0 : tensor<complex<f64>>
+  }
+
+  func.func private @getTensorFilename(index) -> (!Filename)
+
+  //
+  // Main driver that reads matrix from file and calls the sparse kernel.
+  //
+  func.func @entry() {
+    //%d0 = arith.constant 0.0 : complex<f64>
+    %d0 = complex.constant [0.0 : f64, 0.0 : f64] : complex<f64>
+    %c0 = arith.constant 0 : index
+
+    // Setup memory for a single reduction scalar,
+    // initialized to zero.
+    %xdata = memref.alloc() : memref<complex<f64>>
+    memref.store %d0, %xdata[] : memref<complex<f64>>
+    %x = bufferization.to_tensor %xdata : memref<complex<f64>>
+
+    // Read the sparse matrix from file, construct sparse storage.
+    %fileName = call @getTensorFilename(%c0) : (index) -> (!Filename)
+    %a = sparse_tensor.new %fileName : !Filename to tensor<?x?xcomplex<f64>, #SparseMatrix>
+
+    // Call the kernel.
+    %0 = call @kernel_sum_reduce(%a, %x)
+      : (tensor<?x?xcomplex<f64>, #SparseMatrix>, tensor<complex<f64>>) -> tensor<complex<f64>>
+
+    // Print the result for verification.
+    //
+    // CHECK: 30.2
+    // CHECK-NEXT: 22.2
+    //
+    %m = bufferization.to_memref %0 : memref<complex<f64>>
+    %v = memref.load %m[] : memref<complex<f64>>
+    %real = complex.re %v : complex<f64>
+    %imag = complex.im %v : complex<f64>
+    vector.print %real : f64
+    vector.print %imag : f64
+
+    // Release the resources.
+    memref.dealloc %xdata : memref<complex<f64>>
+    sparse_tensor.release %a : tensor<?x?xcomplex<f64>, #SparseMatrix>
+
+    return
+  }
+}

mlir/test/Integration/data/test_symmetric_complex.mtx

+%%MatrixMarket matrix coordinate complex symmetric
+%
+% This is a test sparse matrix in Matrix Market Exchange Format.
+% see https://math.nist.gov/MatrixMarket
+%
+5 5 7
+1 1 5.0 1.0
+1 3 4.1 2.1
+2 2 3.0 3.0
+2 4 2.0 4.0
+3 3 1.0 3.0
+4 4 4.0 2.0
+5 5 5.0 1.0

mlir/test/Integration/data/wide.mtx

-%%MatrixMarket matrix coordinate real general
+%%MatrixMarket matrix coordinate integer general
 %
 % This is a test sparse matrix in Matrix Market Exchange Format.
 % see https://math.nist.gov/MatrixMarket
 %
 4 256 17
 1 1    -1
 1 127   2
 1 128  -3