[Mlir-commits] [mlir] 0271041 - Accept symmetric sparse matrix in Matrix Market Exchange Format.

[Bixia Zheng]

Author: Bixia Zheng
Date: 2021-11-23T19:53:17-08:00
New Revision: 02710413a32afbd6eeeff0fd56c960f4c4e21629

URL: https://github.com/llvm/llvm-project/commit/02710413a32afbd6eeeff0fd56c960f4c4e21629
DIFF: https://github.com/llvm/llvm-project/commit/02710413a32afbd6eeeff0fd56c960f4c4e21629.diff

LOG: Accept symmetric sparse matrix in Matrix Market Exchange Format.

Reviewed By: aartbik

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

Added: 
    mlir/test/Integration/data/test_symmetric.mtx

Modified: 
    mlir/lib/ExecutionEngine/SparseTensorUtils.cpp
    mlir/test/CMakeLists.txt
    mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sum.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/ExecutionEngine/SparseTensorUtils.cpp b/mlir/lib/ExecutionEngine/SparseTensorUtils.cpp
index 75664ecabd07..045bca25c3a5 100644
--- 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 7c5eceebfcc4..f0e1faf8e8b3 100644
--- 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 01583f92e511..230b54567536 100644
--- 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
index 000000000000..71540d05fc55
--- /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