Accept symmetric sparse matrix in Matrix Market Exchange Format.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D114402

diff --git a/mlir/lib/ExecutionEngine/SparseTensorUtils.cpp b/mlir/lib/ExecutionEngine/SparseTensorUtils.cpp
index 75664ec..045bca2 100644 (file)
--- a/mlir/lib/ExecutionEngine/SparseTensorUtils.cpp
+++ b/mlir/lib/ExecutionEngine/SparseTensorUtils.cpp
@@ -496,7 +496,8 @@ static char *toLower(char *token) {
 }
 
 /// Read the MME header of a general sparse matrix of type real.
-static void readMMEHeader(FILE *file, char *name, uint64_t *idata) {
+static void readMMEHeader(FILE *file, char *name, uint64_t *idata,
+                          bool *is_symmetric) {
   char line[1025];
   char header[64];
   char object[64];
@@ -509,11 +510,12 @@ static void readMMEHeader(FILE *file, char *name, uint64_t *idata) {
     fprintf(stderr, "Corrupt header in %s\n", name);
     exit(1);
   }
+  *is_symmetric = (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") ||
-      strcmp(toLower(symmetry), "general")) {
+      (strcmp(toLower(symmetry), "general") && !(*is_symmetric))) {
     fprintf(stderr,
             "Cannot find a general sparse matrix with type real in %s\n", name);
     exit(1);
@@ -579,8 +581,9 @@ static SparseTensorCOO<V> *openSparseTensorCOO(char *filename, uint64_t rank,
   }
   // Perform some file format dependent set up.
   uint64_t idata[512];
+  bool is_symmetric = false;
   if (strstr(filename, ".mtx")) {
-    readMMEHeader(file, filename, idata);
+    readMMEHeader(file, filename, idata, &is_symmetric);
   } else if (strstr(filename, ".tns")) {
     readExtFROSTTHeader(file, filename, idata);
   } else {
@@ -616,6 +619,11 @@ static SparseTensorCOO<V> *openSparseTensorCOO(char *filename, uint64_t rank,
       exit(1);
     }
     tensor->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 && indices[0] != indices[1])
+      tensor->add({indices[1], indices[0]}, value);
   }
   // Close the file and return tensor.
   fclose(file);

diff --git a/mlir/test/CMakeLists.txt b/mlir/test/CMakeLists.txt
index 7c5ecee..f0e1faf 100644 (file)
--- a/mlir/test/CMakeLists.txt
+++ b/mlir/test/CMakeLists.txt
@@ -32,6 +32,7 @@ if (MLIR_INCLUDE_INTEGRATION_TESTS)
   # 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.tns
             ${CMAKE_CURRENT_SOURCE_DIR}/Integration/data/wide.mtx
             ${CMAKE_CURRENT_SOURCE_DIR}/Integration/data/zero.mtx

diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sum.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sum.mlir
index 01583f9..230b545 100644 (file)
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sum.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sum.mlir
@@ -4,7 +4,7 @@
 // RUN:   --func-bufferize --tensor-constant-bufferize --tensor-bufferize \
 // RUN:   --std-bufferize --finalizing-bufferize --lower-affine \
 // RUN:   --convert-vector-to-llvm --convert-memref-to-llvm --convert-std-to-llvm --reconcile-unrealized-casts | \
-// RUN: TENSOR0="%mlir_integration_test_dir/data/test.mtx" \
+// RUN: TENSOR0="%mlir_integration_test_dir/data/test_symmetric.mtx" \
 // RUN: mlir-cpu-runner \
 // RUN:  -e entry -entry-point-result=void  \
 // RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
@@ -18,7 +18,7 @@
 // RUN:   --func-bufferize --tensor-constant-bufferize --tensor-bufferize \
 // RUN:   --std-bufferize --finalizing-bufferize --lower-affine \
 // RUN:   --convert-vector-to-llvm --convert-memref-to-llvm --convert-std-to-llvm --reconcile-unrealized-casts | \
-// RUN: TENSOR0="%mlir_integration_test_dir/data/test.mtx" \
+// RUN: TENSOR0="%mlir_integration_test_dir/data/test_symmetric.mtx" \
 // RUN: mlir-cpu-runner \
 // RUN:  -e entry -entry-point-result=void  \
 // RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
@@ -85,7 +85,7 @@ module {
 
     // Print the result for verification.
     //
-    // CHECK: 28.2
+    // CHECK: 30.2
     //
     %m = memref.buffer_cast %0 : memref<f64>
     %v = memref.load %m[] : memref<f64>

diff --git a/mlir/test/Integration/data/test_symmetric.mtx b/mlir/test/Integration/data/test_symmetric.mtx
new file mode 100644 (file)
index 0000000..71540d0
--- /dev/null
+++ b/mlir/test/Integration/data/test_symmetric.mtx
@@ -0,0 +1,13 @@
+%%MatrixMarket matrix coordinate real 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 3 4.1
+2 2 3.0
+2 4 2.0
+3 3 1.0
+4 4 4.0
+5 5 5.0