[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

1262–1263

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

1321

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.

1358

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)

1374–1375

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

1358

Make the functions inline functions.

[bixia]

Fix a typo.

[aartbik]

mlir/lib/ExecutionEngine/SparseTensorUtils.cpp

1226

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

1252

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

1226

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

1358

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

1298

Maybe add a comment, e.g.

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

1365

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

1365

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

public:
   /// Attempts to read a line from the file.
   char *readLine() {
     if (fgets(line, kColWidth, file))
       return line;
     FATAL("Cannot read next line of %s\n", filename);
   }
 
   /// Reads and parses the file's header.
-  void readHeader() {
+  void readHeader(bool checkValueType = false,
+                  PrimaryType valTp = PrimaryType::kF64) {
     assert(file && "Attempt to readHeader() before openFile()");
     if (strstr(filename, ".mtx"))
-      readMMEHeader();
+      readMMEHeader(checkValueType, valTp);
     else if (strstr(filename, ".tns"))
       readExtFROSTTHeader();
     else
       FATAL("Unknown format %s\n", filename);
     assert(isValid && "Failed to read the header");
   }
 
   /// Gets the MME "pattern" property setting.  Is only valid after
public:
   void assertMatchesShape(uint64_t rank, const uint64_t *shape) const {
     assert(rank == getRank() && "Rank mismatch");
     for (uint64_t r = 0; r < rank; r++)
       assert((shape[r] == 0 || shape[r] == idata[2 + r]) &&
              "Dimension size mismatch");
   }
 
 private:
-  void readMMEHeader();
+  void readMMEHeader(bool checkValueType, PrimaryType valTp);
   void readExtFROSTTHeader();
 
   const char *filename;
   FILE *file = nullptr;
   bool isValid = false;
   bool isPattern_ = false;
   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() {
+void SparseTensorFile::readMMEHeader(bool checkValueType, PrimaryType valTp) {
   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.
+  bool valueKindCompatible = true;
+  if (strcmp(toLower(field), "pattern") == 0)
+    isPattern_ = true;
+  else if (strcmp(toLower(field), "real") == 0)
+    valueKindCompatible =
+        (valTp >= PrimaryType::kF64 && valTp <= PrimaryType::kF32);
+  else if (strcmp(toLower(field), "integer") == 0)
+    valueKindCompatible =
+        (valTp >= PrimaryType::kI64 && valTp <= PrimaryType::kI8);
+  else if (strcmp(toLower(field), "complex") == 0)
+    valueKindCompatible =
+        (valTp >= PrimaryType::kC64 && valTp <= PrimaryType::kC32);
+  else
+    FATAL("Unexpected header field value in %s\n", filename);
+
+  if (checkValueType && !valueKindCompatible)
+    FATAL(
+        "Header field value %s in %s not compatible with tensor element type\n",
+        field, 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.
void SparseTensorFile::readExtFROSTTHeader() {
   // 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;
 }
 
+// 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();
+  stfile.readHeader(/* checkValueType = */ true, valTp);
   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, valTp);      \
       } 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)                                         \
if (!env)
     FATAL("Environment variable %s is not set\n", var);
   return env;
 }
 
 void readSparseTensorShape(char *filename, std::vector<uint64_t> *out) {
   assert(out && "Received nullptr for out-parameter");
   SparseTensorFile stfile(filename);
   stfile.openFile();
-  stfile.readHeader();
+  stfile.readHeader(/* checkValueType = */ false);
   stfile.closeFile();
   const uint64_t rank = stfile.getRank();
   const uint64_t *dimSizes = stfile.getDimSizes();
   out->reserve(rank);
   out->assign(dimSizes, dimSizes + rank);
 }
 
 // TODO: generalize beyond 64-bit indices.

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
+            ${CMAKE_CURRENT_SOURCE_DIR}/Integration/data/wide_integer.mtx
           DESTINATION ${MLIR_INTEGRATION_TEST_DIR}/data/)
 endif()
 
 llvm_canonicalize_cmake_booleans(
   LLVM_BUILD_EXAMPLES
   MLIR_ENABLE_BINDINGS_PYTHON
   MLIR_ENABLE_CUDA_CONVERSIONS
   MLIR_ENABLE_CUDA_RUNNER

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

 // RUN: mlir-opt %s --sparse-compiler | \
-// RUN: TENSOR0="%mlir_integration_test_dir/data/wide.mtx" \
+// RUN: TENSOR0="%mlir_integration_test_dir/data/wide_integer.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
 //
 // Do the same run, but now with SIMDization as well. This should not change the outcome.
 //
 // RUN: mlir-opt %s \
 // RUN:   --sparse-compiler="vectorization-strategy=2 vl=16 enable-simd-index32" | \
-// RUN: TENSOR0="%mlir_integration_test_dir/data/wide.mtx" \
+// RUN: TENSOR0="%mlir_integration_test_dir/data/wide_integer.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<{

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_integer.mtx

+%%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
+1 255   4
+2 2    -5
+2 254   6
+3 3    -7
+4 1     8
+4 2    -9
+4 4    10
+4 99  -11
+4 127  12
+4 128 -13
+4 129  14
+4 250 -15
+4 254  16
+4 256 -17